class_docstr: Gradient wrt x of the Softmax Op.
class_docstr_tok: ["Gradient","wrt","x","of","the","Softmax","Op","."]
code: def c_code(self, node, name, inp, out, sub):
      dy, sm = inp
      (dx,) = out
      axis = self.axis if self.axis is not None else np.MAXDIMS
      fail = sub["fail"]
  
      return dedent(
          f"""
              PyArrayObject* op[3];
              npy_uint32 op_flags[3];
              npy_uint32 iter_flags;
              NpyIter* iter;
              NpyIter_IterNextFunc* get_next;
              char** data_ptr;
  
              int sm_ndim = PyArray_NDIM({sm});
              int axis = {axis};
              int iterate_axis = !(axis == NPY_MAXDIMS || sm_ndim == 1);
  
              // Validate inputs
              if ((PyArray_TYPE({dy}) != NPY_DOUBLE) &&
                  (PyArray_TYPE({dy}) != NPY_FLOAT))
              {{
                  PyErr_SetString(PyExc_TypeError, "types should be float or float64");
                  {fail};
              }}
              if ((PyArray_TYPE({sm}) != NPY_DOUBLE) &&
                  (PyArray_TYPE({sm}) != NPY_FLOAT))
              {{
                  PyErr_SetString(PyExc_TypeError, "types should be float or float64");
                  {fail};
              }}
  
              if (axis < 0) axis = sm_ndim + axis;
              if ((axis < 0) || (iterate_axis && (axis > sm_ndim)))
              {{
                  PyErr_SetString(PyExc_ValueError, "invalid axis in SoftmaxGrad");
                  {fail};
              }}
  
              if (({dx} == NULL)
                  || !(PyArray_CompareLists(PyArray_DIMS({dx}), PyArray_DIMS({sm}), sm_ndim)))
              {{
                  Py_XDECREF({dx});
                  {dx} = (PyArrayObject*)PyArray_SimpleNew(sm_ndim,
                                                           PyArray_DIMS({sm}),
                                                           PyArray_TYPE({sm}));
                  if (!{dx})
                  {{
                      PyErr_SetString(PyExc_MemoryError, "failed to alloc SoftMaxGrad dx output");
                      {fail};
                  }}
              }}
  
              // Create numpy iterator
              op[0] = {dy};
              op[1] = {sm};
              op[2] = {dx};
              op_flags[0] = NPY_ITER_READONLY;
              op_flags[1] = NPY_ITER_READONLY;
              op_flags[2] = NPY_ITER_READWRITE;
              iter_flags = (iterate_axis)? NPY_ITER_MULTI_INDEX : 0;
              iter = NpyIter_MultiNew(
                  3,
                  op,
                  iter_flags,
                  NPY_KEEPORDER,
                  NPY_NO_CASTING,
                  op_flags,
                  NULL
              );
  
              if (iter == NULL)
              {{
                  PyErr_SetString(PyExc_MemoryError, "failed to create softmax iterator");
                  {fail};
              }}
  
              // SoftmaxGrad is applied across the entire array
              if (!iterate_axis)
              {{
                  get_next = NpyIter_GetIterNext(iter, NULL);
                  if (get_next == NULL)
                  {{
                      NpyIter_Deallocate(iter);
                      PyErr_SetString(PyExc_RuntimeError, "Failed to obtain SoftMaxGrad GetIterNext");
                      {fail};
                  }}
                  data_ptr = NpyIter_GetDataPtrArray(iter);
  
                  // Compute and accumulate dy * sm
                  dtype_{dx} sum_dy_times_sm = 0.0;
                  do
                  {{
                      dtype_{dy}* dy_ptr = (dtype_{dy}*)data_ptr[0];
                      dtype_{sm}* sm_ptr = (dtype_{sm}*)data_ptr[1];
                      dtype_{dx}* dx_ptr = (dtype_{dx}*)data_ptr[2];
  
                      *dx_ptr = (dtype_{dx})((*dy_ptr) * (*sm_ptr));
                      sum_dy_times_sm += *dx_ptr;
                  }} while(get_next(iter));
  
                  // Reset Iterator
                  if (NpyIter_GotoIterIndex(iter, 0) == NPY_FAIL)
                  {{
                      PyErr_SetString(PyExc_RuntimeError, "Failed to reset softmax iterator");
                      {fail};
                  }}
  
                  // Subtract sum(dy*sm) * sm
                  do
                  {{
                      dtype_{sm}* sm_ptr = (dtype_{sm}*)data_ptr[1];
                      dtype_{dx}* dx_ptr = (dtype_{dx}*)data_ptr[2];
                      *dx_ptr -= sum_dy_times_sm * ((dtype_{dx})(*sm_ptr));
                  }} while(get_next(iter));
              }}
  
              // SoftmaxGrad is applied across a specific axis
              else {{
                  // Collect axis strides and remove it from iteration
                  npy_intp axis_size = PyArray_DIM({sm}, axis);
                  npy_intp* axis_stride = NpyIter_GetAxisStrideArray(iter, axis);
                  if  (axis_stride == NULL)
                  {{
                      PyErr_SetString(PyExc_RuntimeError, "Failed to obtain softmax axis strides");
                      {fail};
                  }}
                  npy_intp dy_axis_stride = axis_stride[0] / sizeof(dtype_{dy});
                  npy_intp sm_axis_stride = axis_stride[1] / sizeof(dtype_{sm});
                  npy_intp dx_axis_stride = axis_stride[2] / sizeof(dtype_{dx});
  
                  if (NpyIter_RemoveAxis(iter, axis) == NPY_FAIL)
                  {{
                      PyErr_SetString(PyExc_RuntimeError, "Failed to remove SoftmaxGrad axis from iterator");
                      {fail};
                  }}
  
                  // Iterate over remaining axes
                  get_next = NpyIter_GetIterNext(iter, NULL);
                  if (get_next == NULL)
                  {{
                      NpyIter_Deallocate(iter);
                      PyErr_SetString(PyExc_RuntimeError, "Failed to obtain SoftamGrad GetIterNext");
                      {fail};
                  }}
  
                  data_ptr = NpyIter_GetDataPtrArray(iter);
                  do
                  {{
                      dtype_{dy}* dy_axis = (dtype_{dy}*)data_ptr[0];
                      dtype_{sm}* sm_axis = (dtype_{sm}*)data_ptr[1];
                      dtype_{dx}* dx_axis = (dtype_{dx}*)data_ptr[2];
  
                      // Compute and accumulate dy * sm
                      dtype_{dx} sum_dy_times_sm = 0.0;
                      for (npy_intp i = 0; i < axis_size; i++)
                      {{
                          dx_axis[i * dx_axis_stride] = (dtype_{dx})(dy_axis[i * dy_axis_stride] * sm_axis[i * sm_axis_stride]);
                          sum_dy_times_sm += dx_axis[i * dx_axis_stride];
                      }}
  
                      // Subtract sum(dy*sm) * sm
                      for (npy_intp i = 0; i < axis_size; i++)
                      {{
                          dx_axis[i * dx_axis_stride] -= sum_dy_times_sm * (dtype_{dx})(sm_axis[i * sm_axis_stride]);
                      }}
  
                  }} while(get_next(iter));
              }}
              NpyIter_Deallocate(iter);
              """
      )
  
imports: from textwrap import dedent
  import numpy
  import scipy
  from pytensor.graph.basic import Apply
  from pytensor.graph.replace import _vectorize_node
  from pytensor.link.c.op import COp
  from pytensor.tensor.basic import as_tensor_variable
  from pytensor.tensor.elemwise import get_normalized_batch_axes
  from pytensor.tensor.math import gamma, gammaln, neg, sum
label_desc: Use image node_id 7 for calling the SoftmaxGrad obj's underlying member method code with example usage: obj.c_code(node, name, inp, out, sub) and returns: dedent