Closed alvarobartt closed 1 year ago
Something like:
ParamsLike = Union[Dict[str, jnp.DeviceArray], Dict[str, np.ndarray], FrozenDict, VarCollection]
To avoid defining all the types over and over, and also so that the type-hints are aligned between both safejax.serialize and safejax.deserialize :hugs:
safejax.serialize
safejax.deserialize
Something like:
To avoid defining all the types over and over, and also so that the type-hints are aligned between both
safejax.serializeandsafejax.deserialize:hugs: