class_docstr: SM3 optimizer, as described in https://arxiv.org/abs/1901.11150.
class_docstr_tok: ["SM3","optimizer",",","as","described","in","https",":","\/\/arxiv.org\/abs\/1901.11150","."]
code: def _expanded_shape(self, shape, axis):
      # Replaces a `shape` of [M, N, K] with 1 in all dimensions except for i.
      # For eg: i = 1 returns [1, N, 1].
      rank = len(shape)
      return [1] * axis + [shape[axis]] + [1] * (rank - axis - 1)
  
file_path: trax/trax/optimizers/sm3.py
imports: import enum
  from trax.fastmath import numpy
  from trax.optimizers import base
label_desc: Use image node_id 5 for calling the SM3 obj's underlying member method code with example usage: obj._expanded_shape(shape, axis) and returns: unknown