code: def _register_jagged_to_padded_dense(cls):
      # Approximation of fbgemm.jagged_to_padded_dense_forward
      cls.test_inductor_ops.define(
          "jagged_to_padded_dense(Tensor input, Tensor offsets, SymInt max_seq_len, Scalar pad_value) -> Tensor"
      )
  
      def j2pd_meta(inp, offsets, max_seq_len, pad_value):
          return torch.empty(
              (offsets.shape[0] - 1, max_seq_len, inp.shape[1]),
              device=inp.device,
              dtype=inp.dtype,
          )
  
      def j2pd_cuda(inp, offsets, max_seq_len, pad_value):
          res = torch.full(
              (offsets.shape[0] - 1, max_seq_len, inp.shape[1]),
              pad_value,
              device=inp.device,
              dtype=inp.dtype,
          )
          for b in range(offsets.shape[0] - 1):
              for r in range(offsets[b + 1] - offsets[b]):
                  res[b][r] = inp[offsets[b] + r]
          return res
  
      def j2pd_lowering(inp, offsets, max_seq_len, pad_value):
          offsets_loader = offsets.make_loader()
          inp_loader = inp.make_loader()
          jagged_len = inp.get_size()[0]
          offsets_dtype = offsets.get_dtype()
  
          def inner_fn(index):
              batch_idx, seq_idx, emb_idx = index
  
              begin_idx = ops.indirect_indexing(
                  offsets_loader([batch_idx]),
                  jagged_len + 1,
              )
              end_idx = offsets_loader([batch_idx + 1])
              jagged_idx = begin_idx + seq_idx
  
              return ops.masked(
                  ops.lt(
                      ops.index_expr(jagged_idx, offsets_dtype),
                      end_idx,
                  ),
                  lambda: inp_loader([jagged_idx, emb_idx]),
                  pad_value,
              )
  
          return Pointwise.create(
              device=inp.get_device(),
              dtype=inp.get_dtype(),
              inner_fn=inner_fn,
              ranges=[
                  offsets.get_size()[0] - 1,
                  max_seq_len,
                  inp.get_size()[1],
              ],
          )
  
      register_lowering(
          torch.ops.test_inductor_ops.jagged_to_padded_dense,
          type_promotion_kind=None,
      )(j2pd_lowering)
  
      cls.impl_meta.impl("jagged_to_padded_dense", j2pd_meta)
      cls.impl_cuda.impl("jagged_to_padded_dense", j2pd_cuda)
  
file_path: pytorch/test/inductor/test_custom_lowering.py
imports: import unittest
  import torch
  from torch._inductor.ir import Pointwise
  from torch._inductor.lowering import register_lowering
  from torch._inductor.virtualized import ops
  from torch.testing._internal.common_utils import TestCase
  from torch.testing._internal.inductor_utils import HAS_CPU, HAS_CUDA
label_desc: Use image node_id 3 for calling the TestCustomLowering obj's underlying member method code with example usage: obj._register_jagged_to_padded_dense(cls) and returns: torch, res, Pointwise, ops