code: def Sum(axis=None, keepdims=False):
      """Returns a layer that computes sums using one tensor axis.
  
      `Sum` uses one tensor axis to form groups of values and replaces each group
      with the sum of that group. The resulting sum values can either remain in
      their own size 1 axis (`keepdims=True`), or that axis can be removed from the
      overall tensor (default `keepdims=False`), lowering the rank of the tensor by
      one.
  
      Args:
        axis: Axis along which values are grouped for computing a sum; if None,
            compute sum over all elements in tensor.
        keepdims: If `True`, keep the resulting size 1 axis as a separate tensor
            axis; else, remove that axis.
      """
      return Fn(
          "Sum", lambda x: jnp.sum(x, axis=axis, keepdims=keepdims)
      )
  
imports: from absl import logging
  import numpy
  import tensorflow
  from trax import fastmath
  from trax.fastmath import numpy
  from trax.layers import base
  from trax.layers import initializers
  from trax.layers.assert_shape import assert_shape
  from trax.layers.base import Fn
label_desc: Use image node_id 10 for calling a global function with example usage: Sum(axis, keepdims) and returns: Fn