furax.tree

def as_promoted_dtype(x: P) -> P:
    """Promotes the data type of the leaves of a pytree to a common data type.

    Args:
        x: The pytree to promote.

    Examples:
        >>> as_promoted_dtype({'a': jnp.ones(2, jnp.float16), 'b': jnp.ones(2, jnp.float32)})
        {'a': Array([1., 1.], dtype=float32), 'b': Array([1., 1.], dtype=float32)}
    """
    leaves = jax.tree.leaves(x)
    promoted_dtype = jnp.result_type(*leaves)
    result: P = jax.tree.map(
        lambda leaf: (
            jax.ShapeDtypeStruct(leaf.shape, promoted_dtype)
            if isinstance(leaf, jax.ShapeDtypeStruct)
            else jnp.astype(leaf, promoted_dtype)
        ),
        x,
    )
    return result

def as_structure(x: P) -> P:
    """Returns the pytree of ShapedDtypeStruct leaves associated with x.

    Args:
        x: The pytree whose structure will be returned.

    Examples:
        >>> as_structure(jnp.ones(10))
        ShapeDtypeStruct(shape=(10,), dtype=float32)

        >>> as_structure({'a': [jnp.zeros((2, 3)), jnp.array(2)]})
        {'a': [ShapeDtypeStruct(shape=(2, 3), dtype=float32),
        ShapeDtypeStruct(shape=(), dtype=int32, weak_type=True)]}
    """
    result: P = jax.eval_shape(lambda _: _, x)
    return result

def nbytes(x: P) -> int:
    """Returns the total byte size of a pytree.

    Args:
        x: The pytree of array-like leaves (or ``ShapeDtypeStruct`` leaves), whose total byte
            size will be computed.

    Examples:
        >>> nbytes(jnp.ones(10))
        40

        >>> nbytes({'a': [jnp.zeros((2, 3)), jnp.array(2)]})
        28
    """
    # Cast each dim to a Python int before prod: some callers' shapes carry fixed-width numpy
    # int32 (e.g. from np.unique-derived probe shapes), which silently overflows on large arrays
    # instead of promoting to Python's arbitrary-precision int.
    sizes = jax.tree.map(lambda leaf: leaf.dtype.itemsize * prod(int(d) for d in leaf.shape), x)
    return sum(jax.tree.leaves(sizes))

def zeros_like(x: P) -> P:
    """Returns a pytrees of zeros with the same structure as x.

    Args:
        x: The pytree of array-like leaves with ``shape`` and ``dtype`` attributes, whose structure
            will be used to construct the output pytree of zeros.

    Examples:
        >>> zeros_like({'a': jnp.ones(2, dtype=jnp.int32)})
        {'a': Array([0, 0], dtype=int32)}

        >>> zeros_like({'a': jax.ShapeDtypeStruct((2,), jnp.int32)})
        {'a': Array([0, 0], dtype=int32)}
    """
    return full_like(x, 0)

def ones_like(x: P) -> P:
    """Returns a pytrees of ones with the same structure as x.

    Args:
        x: The pytree of array-like leaves with ``shape`` and ``dtype`` attributes, whose structure
            will be used to construct the output pytree of ones.

    Examples:
        >>> ones_like({'a': jnp.zeros(2, dtype=jnp.int32)})
        {'a': Array([1, 1], dtype=int32)}

        >>> ones_like({'a': jax.ShapeDtypeStruct((2,), jnp.int32)})
        {'a': Array([1, 1], dtype=int32)}
    """
    return full_like(x, 1)

def full_like(x: P, fill_value: ScalarLike) -> P:
    """Returns a pytrees of a specified value with the same structure as x.

    Args:
        x: The pytree of array-like leaves with ``shape`` and ``dtype`` attributes, whose structure
            will be used to construct the output pytree of the specified value..
        fill_value: The value to fill with.

    Examples:
        >>> full_like({'a': jnp.array(1, jnp.int32), 'b': jnp.array(2, jnp.float32)}, 3)
        {'a': Array(3, dtype=int32), 'b': Array(3., dtype=float32)}

        >>> full_like({'a': jax.ShapeDtypeStruct((2,), jnp.int32),
        ...          'b': jax.ShapeDtypeStruct((), jnp.float32)}, 3)
        {'a': Array([3, 3], dtype=int32), 'b': Array(3., dtype=float32)}
    """
    result: P = jax.tree.map(lambda leaf: jnp.full_like(leaf, fill_value), x)
    return result

def normal_like(x: P, key: Key[Array, '']) -> P:
    """Returns a pytrees of a normal values with the same structure as x.

    Args:
        x: The pytree of array-like leaves with ``shape`` and ``dtype`` attributes, whose structure
            will be used to construct the output pytree of pseudo-random values.
        key: The PRNGKey to use.

    Examples:
        >>> normal_like({'a': jnp.array(1, jnp.float16),
        ...            'b': jnp.array(2, jnp.float32)}, jax.random.PRNGKey(0))
        {'a': Array(-1.34, dtype=float16), 'b': Array(-1.2515389, dtype=float32)}

        >>> normal_like({'a': jax.ShapeDtypeStruct((2,), jnp.float16),
        ...            'b': jax.ShapeDtypeStruct((), jnp.float32)}, jax.random.PRNGKey(0))
        {'a': Array([-1.34  ,  0.1431], dtype=float16), 'b': Array(-1.2515389, dtype=float32)}
    """
    key_leaves = jax.random.split(key, len(jax.tree.leaves(x)))
    keys = jax.tree.unflatten(jax.tree.structure(x), key_leaves)
    result: P = jax.tree.map(
        lambda leaf, key: jax.random.normal(key, leaf.shape, leaf.dtype), x, keys
    )
    return result

def is_leaf(x: Any) -> bool:
    """Returns true if the input is a Pytree leaf."""
    treedef = jax.tree.structure(x)
    return jax.tree_util.treedef_is_leaf(treedef)

def add(a: PyTree[Array], b: PyTree[Array]) -> PyTree[Array]:
    return apply(operator.add, a, b)

def sub(a: PyTree[Array], b: PyTree[Array]) -> PyTree[Array]:
    return apply(operator.sub, a, b)

def mul(a: PyTree[Array], b: PyTree[Array]) -> PyTree[Array]:
    return apply(operator.mul, a, b)

def truediv(a: PyTree[Array], b: PyTree[Array]) -> PyTree[Array]:
    return apply(operator.truediv, a, b)

def power(a: PyTree[Array], b: PyTree[Array]) -> PyTree[Array]:
    return apply(operator.pow, a, b)

def dot(x: PyTree[Num[Array, '...']], y: PyTree[Num[Array, '...']]) -> Num[Array, '']:
    """Scalar product of two Pytrees.

    If one of the leaves is complex, the hermitian scalar product is returned.

    Args:
        x: The first Pytree.
        y: The second Pytree.

    Examples:
        >>> import furax as fx
        >>> x = {'a': jnp.array([1., 2, 3]), 'b': jnp.array([1, 0])}
        >>> y = {'a': jnp.array([2, -1, 1]), 'b': jnp.array([2, 0])}
        >>> fx.tree.dot(x, y)
        Array(5., dtype=float32)
    """
    # No `jnp.vdot`: it lowers to `dot_general`, which under explicit-axis sharding refuses
    # to pick an output sharding when the contracting axis is sharded (raises ShardingTypeError).
    #
    # Instead, simply use the reduction `sum(conj(x) * y)` for which JAX auto-inserts the
    # all-reduce,so the dot is well-defined when sharded.
    #
    # Trade-off: we lose `dot_general`'s fused high-precision accumulation (vdot can pass
    # `precision=HIGHEST`); a plain `jnp.sum` accumulates in the leaf dtype.
    xy = jax.tree.map(lambda a, b: jnp.sum(jnp.conj(a) * b), x, y)
    return sum(jax.tree.leaves(xy), start=jnp.array(0))

def norm(x: PyTree[Num[Array, '...']]) -> Num[Array, '']:
    """Compute the norm of a PyTree.

    The norm is defined as sqrt(dot(x, x)).

    Args:
        x: The Pytree.

    Examples:
        >>> import furax as fx
        >>> x = {'a': jnp.array([3., 4.]), 'b': jnp.array([0.])}
        >>> fx.tree.norm(x)
        Array(5., dtype=float32)
    """
    # `dot(x, x) = sum(|x|^2)` is already real-valued
    # explicit `.real` keeps the result real-dtyped for complex inputs
    return jnp.sqrt(dot(x, x).real)

def matvec(
    outer_treedef: PyTreeDef | PyTree[Any], a: PyTree[Array], x: PyTree[Array]
) -> PyTree[Array]:
    """Generalized matrix-vector operation, where the matrix and the vector are pytrees.

    The structure of the generalized matrix is the tree product of an outer tree structure
    representing the rows and an inner tree structure that represents the columns. A tree product
    of two tree structures is formed by replacing each leaf of the first tree with a copy of
    the second. The leaves of the generalized matrix must be broadcastable when they belong to
    same inner tree (elements of the same row). There is no such requirement for leaves of different
    inner trees (elements of different rows).

    Args:
        outer_treedef: The outer structure of the generalized matrix.
        a: The generalized matrix, i.e. a pytree whose structure follows the tree product of
            an outer and inner tree structures, with leaves given by the elements of the
            generalized matrix.
        x: The generalized vector, with a structure matching the inner tree structure of `a`.

    Returns:
        A pytree with the same structure as the generalized matrix outer tree structure.

    Examples:
        To represent with pytrees the sparse tensor
            [ a11 a12 ]
            [ a21 a22 ] of shape (2, 2, 100) where a11, a12 and a21 are arrays of 100 elements and
        a22 is zero:

        >>> from numpy.testing import assert_array_equal
        >>> a11, a12, a21 = jax.random.normal(jax.random.key(0), (3, 100))
        >>> a22 = 0
        >>> a = {
        ...     'row1': {'col1': a11, 'col2': a12},
        ...     'row2': {'col1': a21, 'col2': a22},
        ... }
        >>> x = {'col1': 1., 'col2': 2.}
        >>> y = matvec({'row1': 0, 'row2': 0}, a, x)
        >>> assert_array_equal(y['row1'], a11 * x['col1'] + a12 * x['col2'])
        >>> assert_array_equal(y['row2'], a21 * x['col1'] + a22 * x['col2'])
    """
    if not isinstance(outer_treedef, PyTreeDef):
        outer_treedef = jax.tree.structure(outer_treedef)
    outer_leaves = outer_treedef.flatten_up_to(a)
    leaves = []
    for outer_leaf in outer_leaves:
        leaf = sum(jax.tree.leaves(mul(outer_leaf, x)))
        leaves.append(leaf)
    return outer_treedef.unflatten(leaves)

def vecmat(
    x: PyTree[Array], outer_treedef: PyTreeDef | PyTree[Any], a: PyTree[Array]
) -> PyTree[Array]:
    """Generalized vector-matrix operation, where the matrix and the vector are pytrees.

    The structure of the generalized matrix is the tree product of an outer tree structure
    representing the rows and an inner tree structure that represents the columns. A tree product
    of two tree structures is formed by replacing each leaf of the first tree with a copy of
    the second. The leaves of the generalized matrix must be broadcastable when they belong to
    same inner tree (elements of the same row). There is no such requirement for leaves of different
    inner trees (elements of different rows).

    Args:
        outer_treedef: The outer structure of the generalized matrix.
        a: The generalized matrix, i.e. a pytree whose structure follows the tree product of
            an outer and inner tree structures, with leaves given by the elements of the
            generalized matrix.
        x: The generalized vector, with a structure matching the outer tree structure of `a`.

    Returns:
        A pytree with the same structure as the generalized matrix inner tree structure.

    Examples:
        To represent with pytrees the sparse tensor
            [ a11 a12 ]
            [ a21 a22 ] of shape (2, 2, 100) where a11, a12 and a21 are arrays of 100 elements and
        a22 is zero:

        >>> from numpy.testing import assert_array_equal
        >>> a11, a12, a21 = jax.random.normal(jax.random.key(0), (3, 100))
        >>> a22 = 0
        >>> a = {
        ...     'row1': {'col1': a11, 'col2': a12},
        ...     'row2': {'col1': a21, 'col2': a22},
        ... }
        >>> x = {'row1': 1., 'row2': 2.}
        >>> y = vecmat(x, {'row1': 0, 'row2': 0}, a)
        >>> assert_array_equal(y['col1'], a11 * x['row1'] + a21 * x['row2'])
        >>> assert_array_equal(y['col2'], a12 * x['row1'] + a22 * x['row2'])
    """
    if not isinstance(outer_treedef, PyTreeDef):
        outer_treedef = jax.tree.structure(outer_treedef)
    inner_treedef = jax.tree.structure(outer_treedef.flatten_up_to(a)[0])
    transposed_a = jax.tree.transpose(outer_treedef, inner_treedef, a)
    return matvec(inner_treedef, transposed_a, x)

def matmat(
    a_outer_treedef: PyTreeDef | PyTree[Any],
    a: PyTree[Array],
    b_outer_treedef: PyTreeDef | PyTree[Any],
    b: PyTree[Array],
) -> PyTree[Array]:
    """Generalized matrix-matrix operation, where the matrices are pytrees.

    The structure of the generalized matrices is the tree product of an outer tree structure
    representing the rows and an inner tree structure that represents the columns. A tree product
    of two tree structures is formed by replacing each leaf of the first tree with a copy of
    the second. The leaves of the generalized matrix must be broadcastable when they belong to
    same inner tree (elements of the same row). There is no such requirement for leaves of different
    inner trees (elements of different rows).

    Args:
        a_outer_treedef: The outer structure of the first generalized matrix.
        a: The first generalized matrix, i.e. a pytree whose structure follows the tree product of
            an outer and inner tree structures, with leaves given by the elements of the
            generalized matrix.
        b_outer_treedef: The outer structure of the second generalized matrix.
        b: The second generalized matrix.

    Returns:
        A pytree whose structure is the tree product of the outer structure of `a` and the inner
        structure of `b`.
    """
    if not isinstance(a_outer_treedef, PyTreeDef):
        a_outer_treedef = jax.tree.structure(a_outer_treedef)
    if not isinstance(b_outer_treedef, PyTreeDef):
        b_outer_treedef = jax.tree.structure(b_outer_treedef)
    a_outer_leaves = a_outer_treedef.flatten_up_to(a)
    leaves = []
    for a_outer_leaf in a_outer_leaves:
        leaf = vecmat(a_outer_leaf, b_outer_treedef, b)
        leaves.append(leaf)
    return a_outer_treedef.unflatten(leaves)