jjinx 0.0.1__py3-none-any.whl

jinx/errors.py

@@ -0,0 +1,41 @@
+"""J errors for incorrect usage of J primitives."""
+
+
+class BaseJError(Exception):
+    pass
+
+
+class LengthError(BaseJError):
+    pass
+
+
+class DomainError(BaseJError):
+    pass
+
+
+class ValenceError(BaseJError):
+    pass
+
+
+class JIndexError(BaseJError):
+    pass
+
+
+class SpellingError(BaseJError):
+    pass
+
+
+class StackError(BaseJError):
+    pass
+
+
+class EvaluationError(BaseJError):
+    pass
+
+
+class JSyntaxError(BaseJError):
+    pass
+
+
+class JinxNotImplementedError(BaseJError):
+    """Raised when a feature is not implemented in Jinx."""

jinx/execution/executor.py

@@ -0,0 +1,60 @@
+from dataclasses import dataclass
+from typing import Callable
+
+from jinx.vocabulary import Adverb, Conjunction, Noun, Verb
+
+
+@dataclass(frozen=True)
+class Executor[T]:
+    apply_monad: Callable[[Verb[T], Noun[T]], Noun[T]]
+    """Apply monadic form of verb to a noun."""
+
+    apply_dyad: Callable[[Verb[T], Noun[T], Noun[T]], Noun[T]]
+    """Apply dyadic form of verb to two nouns."""
+
+    apply_conjunction: Callable[
+        [Verb[T] | Noun[T], Conjunction, Verb[T]], Verb[T] | Noun[T]
+    ]
+    """Apply conjunction to left and right arguments."""
+
+    apply_adverb: Callable[[Verb[T] | Noun[T], Adverb], Verb[T] | Noun[T]]
+    """Apply adverb to left argument."""
+
+    build_fork: Callable[[Noun[T] | Verb[T], Verb[T], Verb[T]], Verb[T]]
+    """Build fork."""
+
+    build_hook: Callable[[Verb[T], Verb[T]], Verb[T]]
+    """Build hook."""
+
+    ensure_noun_implementation: Callable[[Noun[T]], None]
+    """Ensure that the noun has an implementation."""
+
+    primitive_verb_map: dict[
+        str, tuple[Callable[[T], T] | None, Callable[[T, T], T] | None]
+    ]
+    """Map of primitive verb names to implementations of monad and dyad functions."""
+
+    primitive_adverb_map: dict[str, Callable[[Verb[T]], Verb[T]]]
+    """Map of primitive adverb names to implementation function."""
+
+    primitive_conjuction_map: dict[
+        str, Callable[[Verb[T] | Noun[T], Verb[T] | Noun[T]], Verb[T]]
+    ]
+    """Map of primitive conjunction names to implementation function."""
+
+    noun_to_string: Callable[[Noun[T]], str]
+    """Convert a noun to a string representation for printing."""
+
+
+def load_executor(name: str) -> Executor:
+    if name == "numpy":
+        from jinx.execution.numpy import executor as numpy_executor
+
+        return numpy_executor
+
+    if name == "jax":
+        from jinx.execution.jax import executor as jax_executor
+
+        return jax_executor
+
+    raise ValueError(f"Unknown executor: {name}")

jinx/execution/jax/__init__.py

@@ -0,0 +1,45 @@
+import jax
+from jinx.errors import JinxNotImplementedError
+from jinx.execution.executor import Executor
+from jinx.execution.jax.adverbs import ADVERB_MAP
+from jinx.execution.jax.application import (
+    apply_adverb,
+    #     apply_conjunction,
+    apply_dyad,
+    apply_monad,
+    build_fork,
+    #     build_hook,
+    ensure_noun_implementation,
+)
+from jinx.execution.jax.verbs import VERB_MAP
+
+# from jinx.execution.numpy.conjunctions import CONJUNCTION_MAP
+# from jinx.execution.numpy.conversion import ensure_noun_implementation
+from jinx.execution.numpy.printing import noun_to_string
+
+jax.config.update("jax_dynamic_shapes", True)
+
+
+def make_not_implemented(name: str):
+    def _not_implemented(*_, **__):
+        raise JinxNotImplementedError(
+            f"{name}: not yet implemented in the JAX executor."
+        )
+
+    return _not_implemented
+
+
+executor = Executor[jax.Array](
+    apply_monad=apply_monad,
+    apply_dyad=apply_dyad,
+    apply_conjunction=make_not_implemented("conjunction"),
+    apply_adverb=apply_adverb,
+    build_fork=build_fork,
+    build_hook=make_not_implemented("hook"),
+    ensure_noun_implementation=ensure_noun_implementation,
+    primitive_verb_map=VERB_MAP,
+    primitive_adverb_map=ADVERB_MAP,
+    primitive_conjuction_map={},
+    # Just use the NumPy implementation of printing.
+    noun_to_string=noun_to_string,  # type: ignore[arg-type]
+)

jinx/execution/jax/adverbs.py

@@ -0,0 +1,91 @@
+"""Methods implementing J adverbs."""
+
+import functools
+from typing import Callable
+
+import jax
+import jax.numpy as jnp
+from jinx.errors import JinxNotImplementedError, ValenceError
+from jinx.execution.jax.application import _apply_dyad
+from jinx.execution.numpy.helpers import (
+    maybe_parenthesise_verb_spelling,
+)
+from jinx.vocabulary import Dyad, Monad, Verb
+
+INFINITY = float("inf")
+
+
+def slash_adverb(verb: Verb[jax.Array]) -> Verb[jax.Array]:
+    if verb.dyad is None or verb.dyad.function is None:
+        # Note: this differs from J which still allows the adverb to be applied
+        # to a verb, but may raise an error when the new verb is applied to a noun
+        # and the verb has no dyadic valence.
+        raise ValenceError(f"Verb {verb.spelling} has no dyadic valence.")
+
+    if isinstance(verb.dyad.function, jnp.ufunc) and verb.dyad.is_commutative:
+        monad = verb.dyad.function.reduce
+        dyad = verb.dyad.function.outer
+
+    else:
+        # Slow path: dyad is not a ufunc.
+        # The function is either callable, in which cases it is applied directly,
+        # or a Verb object that needs to be applied indirectly with _apply_dyad().
+        if isinstance(verb.dyad.function, Verb):
+            func = functools.partial(_apply_dyad, verb)  # type: ignore[assignment]
+        else:
+            func = verb.dyad.function  # type: ignore[assignment]
+
+        def _dyad_arg_swap(x: jax.Array, y: jax.Array) -> jax.Array:
+            return func(y, x)
+
+        def _reduce(y: jax.Array) -> jax.Array:
+            y = jnp.atleast_1d(y)
+            y = jnp.flip(y, axis=0)
+            return functools.reduce(_dyad_arg_swap, y)
+
+        monad = _reduce  # type: ignore[assignment]
+        dyad = NotImplemented
+
+    spelling = maybe_parenthesise_verb_spelling(verb.spelling)
+    spelling = f"{verb.spelling}/"
+
+    return Verb[jax.Array](
+        name=spelling,
+        spelling=spelling,
+        monad=Monad(name=spelling, rank=INFINITY, function=monad),
+        dyad=Dyad(
+            name=spelling, left_rank=INFINITY, right_rank=INFINITY, function=dyad
+        ),
+    )
+
+
+def bslash_adverb(verb: Verb[jax.Array]) -> Verb[jax.Array]:
+    # Common cases that have a straightforward optimisation.
+    SPECIAL_MONAD = {
+        "+/": jnp.cumulative_sum,
+        "*/": jnp.cumulative_prod,
+    }
+
+    if verb.spelling in SPECIAL_MONAD:
+        monad_ = SPECIAL_MONAD[verb.spelling]
+
+    else:
+        raise JinxNotImplementedError(
+            f"Adverb \\ applied to verb {verb.spelling} is not yet implemented."
+        )
+
+    spelling = maybe_parenthesise_verb_spelling(verb.spelling)
+    spelling = f"{spelling}\\"
+
+    return Verb(
+        name=spelling,
+        spelling=spelling,
+        monad=Monad(name=spelling, rank=INFINITY, function=monad_),
+        dyad=Dyad(name=spelling, left_rank=0, right_rank=INFINITY, function=None),  # type: ignore[arg-type]
+    )
+
+
+ADVERB_MAP: dict[str, Callable[[Verb[jax.Array]], Verb[jax.Array]]] = {
+    "SLASH": slash_adverb,
+    "BSLASH": bslash_adverb,
+}

jinx/execution/jax/application.py

@@ -0,0 +1,231 @@
+import functools
+from dataclasses import dataclass
+
+import jax
+import jax.numpy as jnp
+from jinx.errors import JinxNotImplementedError, ValenceError
+from jinx.vocabulary import Adverb, DataType, Dyad, Monad, Noun, RankT, Verb
+
+DATATYPE_TO_NP_MAP = {
+    # JAX requires support for int64 to be set via config.
+    DataType.Integer: jnp.int32,
+    DataType.Float: jnp.float64,
+}
+
+
+INFINITY = float("inf")
+
+
+def get_rank(verb_rank: RankT, noun_rank: int) -> int:
+    """Get the rank at which to apply the verb to the noun.
+
+    If the verb rank is negative, it means that the verb rank is subtracted
+    from the noun rank, to a minimum of 0.
+    """
+    if verb_rank < 0:
+        return max(0, noun_rank + verb_rank)  # type: ignore[return-value]
+    return min(verb_rank, noun_rank)  # type: ignore[return-value]
+
+
+@dataclass
+class ArrayCells:
+    cell_shape: tuple[int, ...]
+    frame_shape: tuple[int, ...]
+    cells: jax.Array
+
+
+def split_into_cells(arr: jax.Array, rank: int) -> ArrayCells:
+    """
+    Look at the array shape and rank to determine frame and cell shape.
+
+    The trailing `rank` axes define the cell shape and the preceding
+    axes define the frame shape. E.g. for rank=2:
+
+      arr.shape = (n0, n1, n2, n3, n4)
+                   ----------  ------
+                   ^ frame     ^ cell
+
+    If rank=0, the frame shape is the same as the shape and the monad
+    applies to each atom of the array.
+    """
+    if rank == 0:
+        return ArrayCells(cell_shape=(), frame_shape=arr.shape, cells=arr.ravel())
+
+    return ArrayCells(
+        cell_shape=arr.shape[-rank:],
+        frame_shape=arr.shape[:-rank],
+        cells=arr.reshape(-1, *arr.shape[-rank:]),
+    )
+
+
+def infer_data_type(data: jax.Array) -> DataType:
+    dtype = data.dtype
+    if jnp.issubdtype(dtype, jnp.integer) or jnp.issubdtype(dtype, jnp.bool_):
+        return DataType.Integer
+    if jnp.issubdtype(dtype, jnp.floating):
+        return DataType.Float
+
+    raise NotImplementedError(f"Cannot handle JAX dtype: {dtype}")
+
+
+def convert_noun_to_jax_array(noun: Noun[jax.Array]) -> jax.Array:
+    dtype = DATATYPE_TO_NP_MAP[noun.data_type]
+    if len(noun.data) == 1:
+        # A scalar (ndim == 0) is returned for single element arrays.
+        return jnp.array(noun.data[0], dtype=dtype)  # type: ignore[call-overload]
+    return jnp.array(noun.data, dtype=dtype)  # type: ignore[call-overload]
+
+
+def ensure_noun_implementation(noun: Noun[jax.Array]) -> None:
+    if noun.implementation is None:
+        noun.implementation = convert_noun_to_jax_array(noun)
+
+
+def jax_array_to_noun(array: jax.Array) -> Noun[jax.Array]:
+    data_type = infer_data_type(array)
+    return Noun[jax.Array](data_type=data_type, implementation=array)
+
+
+def apply_monad(verb: Verb[jax.Array], noun: Noun[jax.Array]) -> Noun[jax.Array]:
+    result = _apply_monad(verb, noun.implementation)
+    return jax_array_to_noun(result)
+
+
+def _apply_monad(verb: Verb[jax.Array], arr: jax.Array) -> jax.Array:
+    if verb.monad is None or verb.monad.function is None:
+        raise ValenceError(f"Verb {verb.spelling} has no monadic valence.")
+    if verb.monad.function is NotImplemented:
+        raise JinxNotImplementedError(
+            f"Verb {verb.spelling} monad function is not yet implemented in Jinx."
+        )
+
+    if isinstance(verb.monad.function, Verb):
+        function = functools.partial(_apply_monad, verb.monad.function)
+    else:
+        function = verb.monad.function  # type: ignore[assignment]
+
+    rank = get_rank(verb.monad.rank, arr.ndim)
+
+    if rank == 0:
+        return function(arr)
+
+    array_cells = split_into_cells(arr, rank)
+
+    if array_cells.cells.ndim == 1:
+        cells = function(array_cells.cells)
+    else:
+        # TODO: Use jax.vmap instead
+        cells = jnp.asarray([function(cell) for cell in array_cells.cells])
+
+    # No filling/padding for now...
+    return jnp.asarray(cells).reshape(array_cells.frame_shape + cells[0].shape)
+
+
+def apply_dyad(
+    verb: Verb[jax.Array], noun_1: Noun[jax.Array], noun_2: Noun[jax.Array]
+) -> Noun[jax.Array]:
+    result = _apply_dyad(verb, noun_1.implementation, noun_2.implementation)
+    return jax_array_to_noun(result)
+
+
+def _apply_dyad(
+    verb: Verb[jax.Array], left_arr: jax.Array, right_arr: jax.Array
+) -> jax.Array:
+    if verb.dyad is None or verb.dyad.function is None:
+        raise ValenceError(f"Verb {verb.spelling} has no dyadic valence.")
+    if verb.dyad.function is NotImplemented:
+        raise JinxNotImplementedError(
+            f"Verb {verb.spelling} dyad function is not yet implemented."
+        )
+
+    if isinstance(verb.dyad.function, Verb):
+        function = functools.partial(_apply_dyad, verb.dyad.function)
+    else:
+        function = verb.dyad.function  # type: ignore[assignment]
+
+    left_rank = get_rank(verb.dyad.left_rank, left_arr.ndim)
+    right_rank = get_rank(verb.dyad.right_rank, right_arr.ndim)
+
+    if left_rank == right_rank == 0 and (left_arr.ndim == 0 or right_arr.ndim == 0):
+        return function(left_arr, right_arr)
+
+    left = split_into_cells(left_arr, left_rank)
+    right = split_into_cells(right_arr, right_rank)
+
+    if left.frame_shape == right.frame_shape:
+        cells = [
+            function(left_cell, right_cell)
+            for left_cell, right_cell in zip(left.cells, right.cells, strict=True)
+        ]
+        return jnp.asarray(cells).reshape(left.frame_shape + cells[0].shape)
+
+    raise JinxNotImplementedError(
+        "Dyadic verbs with non-zero rank and different frame shape are not yet implemented."
+    )
+
+
+def apply_adverb(verb_or_noun: Verb | Noun, adverb: Adverb) -> Verb:
+    return adverb.function(verb_or_noun)
+
+
+def build_fork(
+    f: Verb[jax.Array] | Noun[jax.Array], g: Verb[jax.Array], h: Verb[jax.Array]
+) -> Verb[jax.Array]:
+    """Build a fork given verbs f, g, h.
+
+      (f g h) y  ->    (f y) g   (h y)
+    x (f g h) y  ->  (x f y) g (x h y)
+
+    The new verb has infinite rank.
+
+    Note that f can be a noun, in which case there is one fewer function calls.
+    """
+
+    def _monad(y: jax.Array) -> jax.Array:
+        if isinstance(f, Verb) and f.spelling == "[:":
+            hy = _apply_monad(h, y)
+            return _apply_monad(g, hy)
+
+        if isinstance(f, Verb):
+            a = _apply_monad(f, y)
+        else:
+            a = f.implementation
+        b = _apply_monad(h, y)
+        return _apply_dyad(g, a, b)
+
+    def _dyad(x: jax.Array, y: jax.Array) -> jax.Array:
+        if isinstance(f, Verb) and f.spelling == "[:":
+            hy = _apply_dyad(h, x, y)
+            return _apply_monad(g, hy)
+
+        if isinstance(f, Verb):
+            a = _apply_dyad(f, x, y)
+        else:
+            a = f.implementation
+        b = _apply_dyad(h, x, y)
+        return _apply_dyad(g, a, b)
+
+    if isinstance(f, Verb):
+        f_spelling = f"({f.spelling})" if " " in f.spelling else f.spelling
+    else:
+        f_spelling = str(f.implementation)
+
+    g_spelling = f"({g.spelling})" if " " in g.spelling else g.spelling
+    h_spelling = f"({h.spelling})" if " " in h.spelling else h.spelling
+    spelling = f"{f_spelling} {g_spelling} {h_spelling}"
+
+    return Verb[jax.Array](
+        spelling=spelling,
+        name=spelling,
+        monad=Monad(
+            name=spelling,
+            rank=INFINITY,
+            function=_monad,
+        ),
+        dyad=Dyad(
+            name=spelling,
+            left_rank=INFINITY,
+            right_rank=INFINITY,
+            function=_dyad,
+        ),
+    )

jinx/execution/jax/verbs.py

@@ -0,0 +1,90 @@
+from typing import Callable
+
+import jax
+import jax.numpy as jnp
+
+MonadT = Callable[[jax.Array], jax.Array]
+DyadT = Callable[[jax.Array, jax.Array], jax.Array]
+
+
+@jax.jit
+def number_monad(y: jax.Array) -> jax.Array:
+    """# monad: count number of items in y."""
+    if jnp.isscalar(y) or y.shape == ():
+        return jnp.array(1)
+    return jnp.array(y.shape[0])
+
+
+@jax.jit
+def percentco_dyad(x: jax.Array, y: jax.Array) -> jax.Array:
+    return jnp.power(y, 1 / x)
+
+
+@jax.jit
+def plusdot_monad(y: jax.Array) -> jax.Array:
+    """+. monad: returns real and imaginary parts of numbers."""
+    return jnp.stack([jnp.real(y), jnp.imag(y)], axis=-1)
+
+
+@jax.jit
+def plusco_monad(y: jax.Array) -> jax.Array:
+    """+: monad: double the values in the array."""
+    return 2 * y
+
+
+@jax.jit
+def minusdot_monad(y: jax.Array) -> jax.Array:
+    """-.: monad: returns 1 - y."""
+    return 1 - y
+
+
+@jax.jit
+def minusco_monad(y: jax.Array) -> jax.Array:
+    """-: monad: halve the values in the array."""
+    return y / 2
+
+
+@jax.jit
+def minusco_dyad(x: jax.Array, y: jax.Array) -> jax.Array:
+    """-: dyad: match, returns true if x and y have same shape and values."""
+    is_equal = jnp.array_equal(x, y, equal_nan=True)
+    return jnp.asarray(is_equal)
+
+
+@jax.jit
+def hatdot_dyad(x: jax.Array, y: jax.Array) -> jax.Array:
+    """^. dyad: logarithm of y to the base x."""
+    return jnp.log(y) / jnp.log(x)
+
+
+@jax.jit
+def ltco_monad(y: jax.Array) -> jax.Array:
+    """<: monad: decrements the array."""
+    return y - 1
+
+
+@jax.jit
+def gtco_monad(y: jax.Array) -> jax.Array:
+    """>: monad: increments the array."""
+    return y + 1
+
+
+VERB_MAP: dict[str, tuple[MonadT | None, DyadT | None]] = {
+    # VERB: (MONAD, DYAD)
+    "PLUS": (jnp.conj, jnp.add),
+    "PLUSDOT": (plusdot_monad, NotImplemented),
+    "PLUSCO": (plusco_monad, NotImplemented),
+    "MINUS": (jnp.negative, jnp.subtract),
+    "MINUSDOT": (minusdot_monad, NotImplemented),
+    "MINUSCO": (minusco_monad, minusco_dyad),
+    "STAR": (jnp.sign, jnp.multiply),
+    "PERCENT": (jnp.reciprocal, jnp.divide),
+    "HAT": (jnp.exp, jnp.power),
+    "HATDOT": (jnp.log, hatdot_dyad),
+    "LTDOT": (jnp.floor, jnp.minimum),
+    "GTDOT": (jnp.ceil, jnp.maximum),
+    "LTCO": (ltco_monad, jnp.less_equal),
+    "GTCO": (gtco_monad, jnp.greater_equal),
+    "NUMBER": (number_monad, NotImplemented),
+    "BAR": (jnp.abs, NotImplemented),
+}

jinx/execution/numpy/__init__.py

@@ -0,0 +1,29 @@
+import numpy as np
+from jinx.execution.executor import Executor
+from jinx.execution.numpy.adverbs import ADVERB_MAP
+from jinx.execution.numpy.application import (
+    apply_adverb,
+    apply_conjunction,
+    apply_dyad,
+    apply_monad,
+    build_fork,
+    build_hook,
+)
+from jinx.execution.numpy.conjunctions import CONJUNCTION_MAP
+from jinx.execution.numpy.conversion import ensure_noun_implementation
+from jinx.execution.numpy.printing import noun_to_string
+from jinx.execution.numpy.verbs import VERB_MAP
+
+executor = Executor[np.ndarray](
+    apply_monad=apply_monad,
+    apply_dyad=apply_dyad,
+    apply_conjunction=apply_conjunction,
+    apply_adverb=apply_adverb,
+    build_fork=build_fork,
+    build_hook=build_hook,
+    ensure_noun_implementation=ensure_noun_implementation,
+    primitive_verb_map=VERB_MAP,
+    primitive_adverb_map=ADVERB_MAP,
+    primitive_conjuction_map=CONJUNCTION_MAP,  # type: ignore[arg-type]
+    noun_to_string=noun_to_string,
+)