louisdesdoigts / zodiax Goto Github PK

from jax import Array
import zodiax as zdx

class Foo(zdx.Base):
    leaf : Array
    unit : str

    def __call__(self, new_unit: str):
        new_value = convert(self.leaf, self.unit, new_unit) # some simple unit conversion function
        return self.set(('leaf', 'unit'), (new_value, new_unit))

param = 'leaf'
@zdx.filter_jit(param)
@zdx.filter_grad(param)
def f(pytree):
    return pytree('nm').leaf**2

import jax.numpy as np
import zodiax as zdx
import matplotlib.pyplot as plt
from jax import Array

print(zdx.__version__)

class Foo(zdx.Base):
    a : Array
    b : float

    def __init__(self, a, b):
        self.a = a
        self.b = b
    
    def model(self):
        return self.a ** self.b

foo = Foo(np.array([1, 2, 3], dtype=float), 2.)

> 0.4.0

cov = zdx.self_covariance_matrix(foo, ['a', 'b'], zdx.poiss_loglike)
print(cov)
cov = zdx.self_covariance_matrix(foo, ['b', 'a'], zdx.poiss_loglike)
print(cov)

> [[ nan  nan  nan  nan]
>  [ nan  nan  nan  nan]
>  [ nan  nan  nan  inf]
>  [ nan  nan  nan -inf]]
> [[-0.70989895  0.08762264  0.56825095  1.2324241 ]
>  [ 0.08762264  0.2811295  -0.07013908 -0.1521178 ]
>  [ 0.5682511  -0.07013909 -0.19037326 -0.9865156 ]
>  [ 1.2324241  -0.15211779 -0.98651546 -1.8828206 ]]

fish = zdx.self_fisher_matrix(foo, ['a', 'b'], zdx.poiss_loglike)
print(fish)
fish = zdx.self_fisher_matrix(foo, ['b', 'a'], zdx.poiss_loglike)
print(fish)

> [[-3.4253058   0.          0.         -0.42278463]
>  [ 0.         -3.780815    0.         -3.0264199 ]
>  [ 0.          0.         -3.895044   -6.762014  ]
>  [ 0.          0.          0.          0.        ]]
> [[-12.805318    -0.42278463  -3.0264199   -6.7620144 ]
>  [ -0.42278463  -3.4253058    0.           0.        ]
>  [ -3.0264199    0.          -3.780815     0.        ]
>  [ -6.762014     0.           0.          -3.895044  ]]

import zodiax
class Foo(zodiax.ExtendedBase):
    bar : np.ndarray

    def __init__(self, bar):
        self.bar = np.array(bar)

foo = Foo([1., 2, 3])
optim, opt_state = foo.get_optimiser('bar', optax.adam(1.))

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[119], line 10
      8 foo = Foo([1., 2, 3])
      9 # optim, opt_state = foo.get_optimiser(['bar'], [optax.adam(2e-8)])
---> 10 optim, opt_state = foo.get_optimiser('bar', optax.adam(2e-8))

File ~/mambaforge/envs/dlux/lib/python3.11/site-packages/zodiax-0.1.1-py3.11.egg/zodiax/base.py:691, in ExtendedBase.get_optimiser(self, paths, optimisers, get_args, pmap)
    688 optim = multi_transform(opt_dict, param_spec)
    690 # Get filtered optimiser
--> 691 opt_state = optim.init(eqx_filter(self, is_array))
    693 return (optim, opt_state) if not get_args \
    694     else (optim, opt_state, self.get_args(paths, pmap))

File ~/mambaforge/envs/dlux/lib/python3.11/site-packages/optax/_src/combine.py:135, in multi_transform.<locals>.init_fn(params)
    130 if not label_set.issubset(transforms.keys()):
    131   raise ValueError('Some parameters have no corresponding transformation.\n'
    132                    f'Parameter labels: {list(sorted(label_set))} \n'
    133                    f'Transforms keys: {list(sorted(transforms.keys()))} \n')
--> 135 inner_states = {
    136     group: wrappers.masked(tx, make_mask(labels, group)).init(params)
    137     for group, tx in transforms.items()
    138 }
    139 return MultiTransformState(inner_states)

File ~/mambaforge/envs/dlux/lib/python3.11/site-packages/optax/_src/combine.py:136, in <dictcomp>(.0)
    130 if not label_set.issubset(transforms.keys()):
    131   raise ValueError('Some parameters have no corresponding transformation.\n'
    132                    f'Parameter labels: {list(sorted(label_set))} \n'
    133                    f'Transforms keys: {list(sorted(transforms.keys()))} \n')
    135 inner_states = {
--> 136     group: wrappers.masked(tx, make_mask(labels, group)).init(params)
    137     for group, tx in transforms.items()
    138 }
    139 return MultiTransformState(inner_states)

File ~/mambaforge/envs/dlux/lib/python3.11/site-packages/optax/_src/wrappers.py:482, in masked.<locals>.init_fn(params)
    480 def init_fn(params):
    481   mask_tree = mask(params) if callable(mask) else mask
--> 482   masked_params = mask_pytree(params, mask_tree)
    483   return MaskedState(inner_state=inner.init(masked_params))

File ~/mambaforge/envs/dlux/lib/python3.11/site-packages/optax/_src/wrappers.py:478, in masked.<locals>.mask_pytree(pytree, mask_tree)
    477 def mask_pytree(pytree, mask_tree):
--> 478   return tree_map(lambda m, p: p if m else MaskedNode(), mask_tree, pytree)

File ~/mambaforge/envs/dlux/lib/python3.11/site-packages/jax/_src/tree_util.py:206, in tree_map(f, tree, is_leaf, *rest)
    173 """Maps a multi-input function over pytree args to produce a new pytree.
    174 
    175 Args:
   (...)
    203   [[5, 7, 9], [6, 1, 2]]
    204 """
    205 leaves, treedef = tree_flatten(tree, is_leaf)
--> 206 all_leaves = [leaves] + [treedef.flatten_up_to(r) for r in rest]
    207 return treedef.unflatten(f(*xs) for xs in zip(*all_leaves))

File ~/mambaforge/envs/dlux/lib/python3.11/site-packages/jax/_src/tree_util.py:206, in <listcomp>(.0)
    173 """Maps a multi-input function over pytree args to produce a new pytree.
    174 
    175 Args:
   (...)
    203   [[5, 7, 9], [6, 1, 2]]
    204 """
    205 leaves, treedef = tree_flatten(tree, is_leaf)
--> 206 all_leaves = [leaves] + [treedef.flatten_up_to(r) for r in rest]
    207 return treedef.unflatten(f(*xs) for xs in zip(*all_leaves))

ValueError: Expected list, got Array([1., 2., 3.], dtype=float64).

import zodiax as zdx

class Foo(zdx.Base):
    value : str = zdx.field(static=True)
    
    def __init__(self):
        self.value = 'foo'

class FooReturned(Foo):
    def set_static(self, key, value):
        object.__setattr__(self, key, value)
        return self

foo = FooReturned()

# This prints the updated object, like we expect an immutable class to work
print(foo.set_static('value', 'bar'))

# This prints the updated object too, even though we never re-assigned the variable 'foo'
print(foo)

FooReturned(value='bar')
FooReturned(value='bar')

class FooNotReturned(Foo):
    def set_static(self, key, value):
        object.__setattr__(self, key, value)

foo = FooNotReturned()
print(foo.set_static('value', 'bar'))

foo = foo.set_static('value', 'bar')
print(foo)

None
None

from functools import partial

class FooException(Foo):
    test : float = 1.

    @partial(jax.jit, static_argnums=(1, 2,))
    def set_static(self, key, value):
        object.__setattr__(self, key, value)
        return self

    def set(self, key, value):
        try:
            return super().set(key, value)
        except AssertionError:
            try:
                return self.set_static(key, value)
            except ValueError:
                raise ValueError(f"Static leaves can only be set to hashable data types. Tried to set the leaf '{key}' with a {type(value)} type.")
        return self

foo = FooException()
print(foo.set('test', 2.))

foo = FooException()
print(foo.set('test', 2.).test)

print(foo.set('value', 'bar'))
print(foo)

foo = foo.set('value', 'bar')
print(foo)

FooException(value='foo', test=f32[])
2.0
FooException(value='bar', test=1.0)
FooException(value='foo', test=1.0)
FooException(value='bar', test=1.0)

foo = FooException()
print(foo.set('value', jax.numpy.ones(2)))

...
ValueError: Static leaves can only be set to hashable data types. Tried to set the leaf 'value' with a <class 'jaxlib.xla_extension.ArrayImpl'> type.

import equinox as eqx
import pytreeclass as pytc
import jax


class Tree(eqx.Module):
    weight: jax.Array = jax.numpy.array([-1, 2, 3])
    bias: jax.Array = jax.numpy.array([1])
    counter: int = 1

    @property
    def at(self):
        return pytc.AtIndexer(self, ())


tree = Tree()

tree = (
    tree.at["counter"]
    .set(1)  # set counter to 1
    .at[jax.tree_map(lambda x: x < 0, tree)]
    .set(0)  # set negative values to 0
    .at["bias"].set(100)  # set bias to 100
)

print(tree.weight)
# [0 2 3]
print(tree.bias)
# 100
print(tree.counter)
# 1

import zodiax as zdx
class Foo(zdx.Base):
    param : float = 1.

foo = Foo()

# Works
bar = foo.set(['param'], [None])

# ValueError
bar = foo.set('param', None)

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Input In [10], in <cell line: 7>()
      5 foo = Foo()
      6 bar = foo.set(['param'], [None])
----> 7 bar = foo.set('param', None)

File ~/mambaforge/envs/dlux/lib/python3.10/site-packages/zodiax/base.py:296, in Base.set(self, paths, values, pmap)
    275 def set(self   : Pytree,
    276         paths  : Union[str, list],
    277         values : Union[Any, list],
    278         pmap   : dict = None) -> Pytree:
    279     """
    280     Set the leaves specified by paths with values.
    281 
   (...)
    294         The pytree with leaves specified by paths updated with values.
    295     """
--> 296     new_paths, new_values = self._format(paths, values, pmap)
    298     # Define 'where' function and update pytree
    299     get_leaves_fn = lambda pytree: pytree._get_leaves(new_paths)

ValueError: not enough values to unpack (expected 2, got 1)