code: def diagonal_or_scalar_step(self, t0, t1, y0, extra0):
      del extra0
      dt = t1 - t0
      rdt = 1 / dt
      sqrt_dt = dt.sqrt()
      I_k, I_k0 = self.bm(t0, t1, return_U=True)
      I_kk = (I_k**2 - dt) * _r2
      I_kkk = (I_k**3 - 3 * dt * I_k) * _r6
  
      y1 = y0
      H0, H1 = [], []
      for s in range(srid2.STAGES):
          H0s, H1s = (
              y0,
              y0,
          )  # Values at the current stage to be accumulated.
          for j in range(s):
              f = self.sde.f(t0 + srid2.C0[j] * dt, H0[j])
              g = self.sde.g(t0 + srid2.C1[j] * dt, H1[j])
              g = g.squeeze(2) if g.dim() == 3 else g
              H0s = (
                  H0s
                  + srid2.A0[s][j] * f * dt
                  + srid2.B0[s][j] * g * I_k0 * rdt
              )
              H1s = (
                  H1s
                  + srid2.A1[s][j] * f * dt
                  + srid2.B1[s][j] * g * sqrt_dt
              )
          H0.append(H0s)
          H1.append(H1s)
  
          f = self.sde.f(t0 + srid2.C0[s] * dt, H0s)
          g_weight = (
              srid2.beta1[s] * I_k
              + srid2.beta2[s] * I_kk / sqrt_dt
              + srid2.beta3[s] * I_k0 * rdt
              + srid2.beta4[s] * I_kkk * rdt
          )
          g_prod = self.sde.g_prod(t0 + srid2.C1[s] * dt, H1s, g_weight)
          y1 = y1 + srid2.alpha[s] * f * dt + g_prod
      return y1, ()
  
file_path: torchsde/torchsde/_core/methods/srk.py
imports: from .tableaus import sra1, srid2
  from ..None import adjoint_sde
  from ..None import base_solver
  from ...settings import SDE_TYPES, NOISE_TYPES, LEVY_AREA_APPROXIMATIONS
label_desc: Use image node_id 3 for calling the SRK obj's underlying member method code with example usage: obj.diagonal_or_scalar_step(t0, t1, y0, extra0) and returns: y1,  