Module adseq.implementations.sorted_array

Priority queue. Every insert is a full memory move. Better would be circular but then finding the insertion point is a bit harder

Classes

class SortedArray (buffer: jax.Array)
Expand source code
class SortedArray(typing.NamedTuple):
    buffer: jax.Array
    @classmethod
    def init(cls, delay, capacity=None, grad=False):
        return cls(
                jnp.full(delay if capacity is None else capacity, INT_MAX, 'float32' if grad else 'int32'),
                )
    @classmethod
    def sized(cls, n):
        "wish I could use __class_getitem__"
        return type(f'{cls.__name__}[{n}]',
                    cls.__bases__,
                    {**cls.__dict__,
                     "init": functools.partial(cls.init, capacity=n)})
    def enqueue(self, n):
        return _enqueue(self, n)

    def pop(self, n):
        return _pop(self, n)

SortedArray(buffer,)

Ancestors

  • builtins.tuple

Static methods

def init(delay, capacity=None, grad=False)
def sized(n)

wish I could use class_getitem

Instance variables

var buffer : jax.Array
Expand source code
class SortedArray(typing.NamedTuple):
    buffer: jax.Array
    @classmethod
    def init(cls, delay, capacity=None, grad=False):
        return cls(
                jnp.full(delay if capacity is None else capacity, INT_MAX, 'float32' if grad else 'int32'),
                )
    @classmethod
    def sized(cls, n):
        "wish I could use __class_getitem__"
        return type(f'{cls.__name__}[{n}]',
                    cls.__bases__,
                    {**cls.__dict__,
                     "init": functools.partial(cls.init, capacity=n)})
    def enqueue(self, n):
        return _enqueue(self, n)

    def pop(self, n):
        return _pop(self, n)

Alias for field number 0

Methods

def enqueue(self, n)
Expand source code
def enqueue(self, n):
    return _enqueue(self, n)
def pop(self, n)
Expand source code
def pop(self, n):
    return _pop(self, n)