code: def main(_):
      # Make the network and optimiser.
      net = hk.without_apply_rng(hk.transform(net_fn))
      opt = optax.adam(1e-3)
  
      # Define layerwise sparsities
      def module_matching(s):
          def match_func(m, n, k):
              return m.endswith(s) and not sparsity_ignore(m, n, k)
  
          return match_func
  
      module_sparsity = (
          (module_matching("linear"), 0.98),
          (module_matching("linear_1"), 0.9),
      )
  
      # Training loss (cross-entropy).
      @jax.jit
      def loss(params: hk.Params, batch: Batch) -> jax.Array:
          """Compute the loss of the network, including L2."""
          logits = net.apply(params, batch)
          labels = jax.nn.one_hot(batch["label"], 10)
  
          l2_loss = 0.5 * sum(
              jnp.sum(jnp.square(p))
              for p in jax.tree_util.tree_leaves(params)
          )
          softmax_xent = -jnp.sum(labels * jax.nn.log_softmax(logits))
          softmax_xent /= labels.shape[0]
  
          return softmax_xent + 1e-4 * l2_loss
  
      # Evaluation metric (classification accuracy).
      @jax.jit
      def accuracy(params: hk.Params, batch: Batch) -> jax.Array:
          predictions = net.apply(params, batch)
          return jnp.mean(
              jnp.argmax(predictions, axis=-1) == batch["label"]
          )
  
      @jax.jit
      def get_updates(
          params: hk.Params,
          opt_state: optax.OptState,
          batch: Batch,
      ) -> tuple[hk.Params, optax.OptState]:
          """Learning rule (stochastic gradient descent)."""
          grads = jax.grad(loss)(params, batch)
          updates, opt_state = opt.update(grads, opt_state)
          return updates, opt_state
  
      # We maintain avg_params, the exponential moving average of the "live" params.
      # avg_params is used only for evaluation (cf. https://doi.org/10.1137/0330046)
      @jax.jit
      def ema_update(params, avg_params):
          return optax.incremental_update(
              params, avg_params, step_size=0.001
          )
  
      # Make datasets.
      train = load_dataset("train", shuffle=True, batch_size=1000)
      train_eval = load_dataset(
          "train", shuffle=False, batch_size=10000
      )
      test_eval = load_dataset("test", shuffle=False, batch_size=10000)
  
      # Implemenation note: It is possible to avoid pruned_params and just use
      # a single params which progressively gets pruned.  The updates also don't
      # need to masked in such an implementation.  The current implementation
      # attempts to mimic the way the current TF implementation which allows for
      # previously inactivated connections to become active again if active values
      # drop below their value.
  
      # Initialize network and optimiser; note we draw an input to get shapes.
      pruned_params = params = avg_params = net.init(
          jax.random.PRNGKey(42), next(train)
      )
  
      masks = update_mask(params, 0.0, module_sparsity)
      opt_state = opt.init(params)
  
      # Train/eval loop.
      for step in range(10001):
          if step % 1000 == 0:
              # Periodically evaluate classification accuracy on train & test sets.
              avg_params = apply_mask(
                  avg_params, masks, module_sparsity
              )
              train_accuracy = accuracy(avg_params, next(train_eval))
              test_accuracy = accuracy(avg_params, next(test_eval))
              (
                  total_params,
                  total_nnz,
                  per_layer_sparsities,
              ) = get_sparsity(avg_params)
              (
                  train_accuracy,
                  test_accuracy,
                  total_nnz,
                  per_layer_sparsities,
              ) = jax.device_get(
                  (
                      train_accuracy,
                      test_accuracy,
                      total_nnz,
                      per_layer_sparsities,
                  )
              )
              print(
                  f"[Step {step}] Train / Test accuracy: "
                  f"{train_accuracy:.3f} / {test_accuracy:.3f}."
              )
              print(
                  f"Non-zero params / Total: {total_nnz} / {total_params}; "
                  f"Total Sparsity: {1. - total_nnz / total_params:.3f}"
              )
  
          # Do SGD on a batch of training examples.
          pruned_params = apply_mask(params, masks, module_sparsity)
          updates, opt_state = get_updates(
              pruned_params, opt_state, next(train)
          )
          # applying a straight-through estimator here (that is not masking
          # the updates) leads to much worse performance.
          updates = apply_mask(updates, masks, module_sparsity)
          params = optax.apply_updates(params, updates)
          # we start pruning at iteration 1000 and end at iteration 8000
          progress = min(max((step - 1000.0) / 8000.0, 0.0), 1.0)
          if step % 200 == 0:
              sparsity_fraction = zhugupta_func(progress)
              masks = update_mask(
                  params, sparsity_fraction, module_sparsity
              )
          avg_params = ema_update(params, avg_params)
      print(per_layer_sparsities)
  
imports: from collections.abc import Iterator, Mapping, Sequence
  import functools
  from typing import Callable
  from absl import app
  import haiku
  import jax
  import jax.numpy
  import numpy
  import optax
  import tensorflow_datasets
label_desc: Use image node_id 10 for calling a global function with example usage: main(_) and returns: match_func, unknown, jnp, updates,  opt_state, optax, unknown