Simplify JAX Scalify MNIST examples using jax_scalify.tree methods.

examples/mnist/mnist_classifier_from_scratch.py

@@ -105,7 +105,7 @@ def data_stream():
  params = init_random_params(param_scale, layer_sizes)
  # Transform parameters to `ScaledArray` and proper dtype.
  params = jsa.as_scaled_array(params, scale=scale_dtype(param_scale))
- params = jax.tree_util.tree_map(lambda v: v.astype(training_dtype), params, is_leaf=jsa.core.is_scaled_leaf)
+ params = jsa.tree.astype(params, training_dtype)
 
  @jit
  @jsa.scalify
@@ -119,7 +119,7 @@ def update(params, batch):
  batch = next(batches)
  # Scaled micro-batch + training dtype cast.
  batch = jsa.as_scaled_array(batch, scale=scale_dtype(1))
- batch = jax.tree_util.tree_map(lambda v: v.astype(training_dtype), batch, is_leaf=jsa.core.is_scaled_leaf)
+ batch = jsa.tree.astype(batch, training_dtype)
 
  with jsa.ScalifyConfig(rounding_mode=jsa.Pow2RoundMode.DOWN, scale_dtype=scale_dtype):
  params = update(params, batch)

examples/mnist/mnist_classifier_from_scratch_fp8.py

@@ -22,7 +22,6 @@
 import time
 
 import datasets
-import jax
 import jax.numpy as jnp
 import ml_dtypes
 import numpy as np
@@ -133,7 +132,7 @@ def data_stream():
  params = init_random_params(param_scale, layer_sizes)
  # Transform parameters to `ScaledArray` and proper dtype.
  params = jsa.as_scaled_array(params, scale=scale_dtype(param_scale))
- params = jax.tree_util.tree_map(lambda v: v.astype(training_dtype), params, is_leaf=jsa.core.is_scaled_leaf)
+ params = jsa.tree.astype(params, training_dtype)
 
  @jit
  @jsa.scalify
@@ -147,7 +146,7 @@ def update(params, batch):
  batch = next(batches)
  # Scaled micro-batch + training dtype cast.
  batch = jsa.as_scaled_array(batch, scale=scale_dtype(1))
- batch = jax.tree_util.tree_map(lambda v: v.astype(training_dtype), batch, is_leaf=jsa.core.is_scaled_leaf)
+ batch = jsa.tree.astype(batch, training_dtype)
 
  with jsa.ScalifyConfig(rounding_mode=jsa.Pow2RoundMode.DOWN, scale_dtype=scale_dtype):
  params = update(params, batch)

examples/mnist/mnist_classifier_mlp_flax.py

@@ -79,8 +79,8 @@ def update(model, optimizer, model_state, opt_state, batch):
  key = jax.random.PRNGKey(42)
  use_scalify: bool = True
 
- # training_dtype = np.dtype(np.float16)
  training_dtype = np.dtype(np.float16)
+ optimizer_dtype = np.dtype(np.float16)
  scale_dtype = np.float32
 
  train_images, train_labels, test_images, test_labels = datasets.mnist()
@@ -102,27 +102,24 @@ def data_stream():
  model_state = model.init(key, np.zeros((batch_size, mnist_img_size), dtype=training_dtype))
  # Optimizer & optimizer state.
  # opt = optax.sgd(learning_rate=step_size)
- opt = optax.adam(learning_rate=step_size, eps=1e-5)
+ opt = optax.adam(learning_rate=step_size, eps=2**-16)
  opt_state = opt.init(model_state)
  # Freeze model, optimizer (with step size).
  update_fn = partial(update, model, opt)
 
  if use_scalify:
- # Transform parameters to `ScaledArray` and proper dtype.
+ # Transform parameters to `ScaledArray`.
  model_state = jsa.as_scaled_array(model_state, scale=scale_dtype(1.0))
  opt_state = jsa.as_scaled_array(opt_state, scale=scale_dtype(0.0001))
-
- model_state = jax.tree_util.tree_map(
- lambda v: v.astype(training_dtype), model_state, is_leaf=jsa.core.is_scaled_leaf
- )
  # Scalify the update function as well.
  update_fn = jsa.scalify(update_fn)
- else:
- model_state = jax.tree_util.tree_map(lambda v: v.astype(training_dtype), model_state)
+ # Convert the model state (weights) & optimizer state to proper dtype.
+ model_state = jsa.tree.astype(model_state, training_dtype)
+ opt_state = jsa.tree.astype(opt_state, optimizer_dtype, floating_only=True)
 
  print(f"Using Scalify: {use_scalify}")
  print(f"Training data format: {training_dtype.name}")
- # print(f"Optimizer data format: {training_dtype.name}")
+ print(f"Optimizer data format: {optimizer_dtype.name}")
  print("")
 
  update_fn = jax.jit(update_fn)
@@ -134,7 +131,7 @@ def data_stream():
  for _ in range(num_batches):
  batch = next(batches)
  # Scaled micro-batch + training dtype cast.
- batch = jax.tree_util.tree_map(lambda v: v.astype(training_dtype), batch)
+ batch = jsa.tree.astype(batch, training_dtype)
  if use_scalify:
  batch = jsa.as_scaled_array(batch, scale=scale_dtype(1))
  with jsa.ScalifyConfig(rounding_mode=jsa.Pow2RoundMode.DOWN, scale_dtype=scale_dtype):

jax_scalify/core/datatype.py

@@ -132,10 +132,13 @@ def make_scaled_scalar(val: Array, scale_dtype: Optional[DTypeLike] = None) -> S
 
 
 def is_scaled_leaf(val: Any) -> bool:
- """Is input a JAX PyTree (scaled) leaf, including ScaledArray.
+ """Is input a normal JAX PyTree leaf (i.e. `Array`) or `ScaledArray1.
 
- This function is useful for JAX PyTree handling where the user wants
- to keep the ScaledArray datastructures (i.e. not flattened as a pair of arrays).
+ This function is useful for JAX PyTree handling with `jax.tree` methods where
+ the user wants to keep the ScaledArray data structures (i.e. not flattened as a
+ pair of arrays).
+
+ See `jax_scalify.tree` for PyTree `jax.tree` methods compatible with `ScaledArray`.
  """
  # TODO: check Numpy scalars as well?
  return np.isscalar(val) or isinstance(val, (Array, np.ndarray, ScaledArray))