JAX 中文文档（四）（3）

>>> x = jnp.arange(5, dtype='int8')
>>> 2 * x
Array([0, 2, 4, 6, 8], dtype=int8) 

>>> jnp.int32(2) * x
Array([0, 2, 4, 6, 8], dtype=int32) 

>>> jnp.asarray(2)
Array(2, dtype=int32, weak_type=True) 

>>> jnp.asarray(2, dtype='int32')
Array(2, dtype=int32) 
```## 严格的 dtype 提升
在某些情况下，禁用隐式类型提升行为并要求所有提升都是显式的可能很有用。可以通过在 JAX 中将`jax_numpy_dtype_promotion`标志设置为`'strict'`来实现。在本地，可以通过上下文管理器来完成：
```py
>>> x = jnp.float32(1)
>>> y = jnp.int32(1)
>>> with jax.numpy_dtype_promotion('strict'):
...   z = x + y  
...
Traceback (most recent call last):
TypePromotionError: Input dtypes ('float32', 'int32') have no available implicit
dtype promotion path when jax_numpy_dtype_promotion=strict. Try explicitly casting
inputs to the desired output type, or set jax_numpy_dtype_promotion=standard. 

>>> with jax.numpy_dtype_promotion('strict'):
...   z = x + 1
>>> print(z)
2.0 

jax.config.update('jax_numpy_dtype_promotion', 'strict') 

jax.config.update('jax_numpy_dtype_promotion', 'standard') 

[1, "a", object()]  # 3 leaves
(1, (2, 3), ())  # 3 leaves
[1, {"k1": 2, "k2": (3, 4)}, 5]  # 5 leaves 

(a1, {"k1": a2, "k2": a3}) 

(None, {"k1": None, "k2": 0}) 

(None, 0)  # equivalent to (None, {"k1": 0, "k2": 0}) 

0 

from jax.tree_util import tree_structure
print(tree_structure(object)) 

from jax.tree_util import tree_flatten, tree_unflatten
import jax.numpy as jnp
# The structured value to be transformed
value_structured = [1., (2., 3.)]
# The leaves in value_flat correspond to the `*` markers in value_tree
value_flat, value_tree = tree_flatten(value_structured)
print(f"{value_flat=}\n{value_tree=}")
# Transform the flat value list using an element-wise numeric transformer
transformed_flat = list(map(lambda v: v * 2., value_flat))
print(f"{transformed_flat=}")
# Reconstruct the structured output, using the original
transformed_structured = tree_unflatten(value_tree, transformed_flat)
print(f"{transformed_structured=}") 

value_flat=[1.0, 2.0, 3.0]
value_tree=PyTreeDef([*, (*, *)])
transformed_flat=[2.0, 4.0, 6.0]
transformed_structured=[2.0, (4.0, 6.0)] 

from collections import namedtuple
Point = namedtuple('Point', ['x', 'y'])
example_containers = [
    (1., [2., 3.]),
    (1., {'b': 2., 'a': 3.}),
    1.,
    None,
    jnp.zeros(2),
    Point(1., 2.)
]
def show_example(structured):
  flat, tree = tree_flatten(structured)
  unflattened = tree_unflatten(tree, flat)
  print(f"{structured=}\n  {flat=}\n  {tree=}\n  {unflattened=}")
for structured in example_containers:
  show_example(structured) 

structured=(1.0, [2.0, 3.0])
  flat=[1.0, 2.0, 3.0]
  tree=PyTreeDef((*, [*, *]))
  unflattened=(1.0, [2.0, 3.0])
structured=(1.0, {'b': 2.0, 'a': 3.0})
  flat=[1.0, 3.0, 2.0]
  tree=PyTreeDef((*, {'a': *, 'b': *}))
  unflattened=(1.0, {'a': 3.0, 'b': 2.0})
structured=1.0
  flat=[1.0]
  tree=PyTreeDef(*)
  unflattened=1.0
structured=None
  flat=[]
  tree=PyTreeDef(None)
  unflattened=None
structured=Array([0., 0.], dtype=float32)
  flat=[Array([0., 0.], dtype=float32)]
  tree=PyTreeDef(*)
  unflattened=Array([0., 0.], dtype=float32)
structured=Point(x=1.0, y=2.0)
  flat=[1.0, 2.0]
  tree=PyTreeDef(CustomNode(namedtuple[Point], [*, *]))
  unflattened=Point(x=1.0, y=2.0) 

class Special(object):
  def __init__(self, x, y):
    self.x = x
    self.y = y
  def __repr__(self):
    return "Special(x={}, y={})".format(self.x, self.y)
show_example(Special(1., 2.)) 

structured=Special(x=1.0, y=2.0)
  flat=[Special(x=1.0, y=2.0)]
  tree=PyTreeDef(*)
  unflattened=Special(x=1.0, y=2.0) 

from jax.tree_util import register_pytree_node
class RegisteredSpecial(Special):
  def __repr__(self):
    return "RegisteredSpecial(x={}, y={})".format(self.x, self.y)
def special_flatten(v):
  """Specifies a flattening recipe.
 Params:
 v: the value of registered type to flatten.
 Returns:
 a pair of an iterable with the children to be flattened recursively,
 and some opaque auxiliary data to pass back to the unflattening recipe.
 The auxiliary data is stored in the treedef for use during unflattening.
 The auxiliary data could be used, e.g., for dictionary keys.
 """
  children = (v.x, v.y)
  aux_data = None
  return (children, aux_data)
def special_unflatten(aux_data, children):
  """Specifies an unflattening recipe.
 Params:
 aux_data: the opaque data that was specified during flattening of the
 current treedef.
 children: the unflattened children
 Returns:
 a re-constructed object of the registered type, using the specified
 children and auxiliary data.
 """
  return RegisteredSpecial(*children)
# Global registration
register_pytree_node(
    RegisteredSpecial,
    special_flatten,    # tell JAX what are the children nodes
    special_unflatten   # tell JAX how to pack back into a RegisteredSpecial
)
show_example(RegisteredSpecial(1., 2.)) 

structured=RegisteredSpecial(x=1.0, y=2.0)
  flat=[1.0, 2.0]
  tree=PyTreeDef(CustomNode(RegisteredSpecial[None], [*, *]))
  unflattened=RegisteredSpecial(x=1.0, y=2.0) 

from jax.tree_util import register_pytree_node_class
@register_pytree_node_class
class RegisteredSpecial2(Special):
  def __repr__(self):
    return "RegisteredSpecial2(x={}, y={})".format(self.x, self.y)
  def tree_flatten(self):
    children = (self.x, self.y)
    aux_data = None
    return (children, aux_data)
  @classmethod
  def tree_unflatten(cls, aux_data, children):
    return cls(*children)
show_example(RegisteredSpecial2(1., 2.)) 

structured=RegisteredSpecial2(x=1.0, y=2.0)
  flat=[1.0, 2.0]
  tree=PyTreeDef(CustomNode(RegisteredSpecial2[None], [*, *]))
  unflattened=RegisteredSpecial2(x=1.0, y=2.0) 

class MyTree:
  def __init__(self, a):
    self.a = jnp.asarray(a)
register_pytree_node(MyTree, lambda tree: ((tree.a,), None),
    lambda _, args: MyTree(*args))
tree = MyTree(jnp.arange(5.0))
jax.vmap(lambda x: x)(tree)      # Error because object() is passed to MyTree.
jax.jacobian(lambda x: x)(tree)  # Error because MyTree(...) is passed to MyTree 

class MyTree:
  def __init__(self, a):
    if not (type(a) is object or a is None or isinstance(a, MyTree)):
      a = jnp.asarray(a)
    self.a = a 

def tree_unflatten(aux_data, children):
  del aux_data  # unused in this class
  obj = object.__new__(MyTree)
  obj.a = a
  return obj