parameters: (reg_name)
return_statement: return do_register
docstring: Register a subclass of CustomOpProp to the registry with name reg_name.
docstring_summary: Register a subclass of CustomOpProp to the registry with name reg_name.
docstring_tokens: Register a subclass of CustomOpProp to the registry with name reg_name .
function: def register(reg_name):
      """Register a subclass of CustomOpProp to the registry with name reg_name."""
      def do_register(prop_cls):
          """Register a subclass of CustomOpProp to the registry."""
          fb_functype = CFUNCTYPE(c_int, c_int, POINTER(c_void_p), POINTER(c_int),
                                  POINTER(c_int), c_int, c_void_p)
          del_functype = CFUNCTYPE(c_int, c_void_p)
  
          infershape_functype = CFUNCTYPE(c_int, c_int, POINTER(c_int),
                                          POINTER(POINTER(mx_int)), c_void_p)
          infertype_functype = CFUNCTYPE(c_int, c_int, POINTER(c_int), c_void_p)
          inferstorage_functype = CFUNCTYPE(c_int, c_int, POINTER(c_int), c_void_p)
          inferstorage_backward_functype = CFUNCTYPE(c_int, c_int, POINTER(c_int), \
                                                     POINTER(c_int), c_void_p)
          list_functype = CFUNCTYPE(c_int, POINTER(POINTER(POINTER(c_char))), c_void_p)
          deps_functype = CFUNCTYPE(c_int, c_int_p, c_int_p, c_int_p,
                                    c_int_p, POINTER(c_int_p), c_void_p)
          createop_functype = CFUNCTYPE(c_int, c_char_p, c_int, POINTER(POINTER(mx_uint)),
                                        POINTER(c_int), POINTER(c_int),
                                        POINTER(MXCallbackList), c_void_p)
          req_enum = ('null', 'write', 'inplace', 'add')
  
          def creator(op_type, argc, keys, vals, ret):
              """internal function"""
              assert py_str(op_type) == reg_name
              kwargs = dict([(py_str(keys[i]), py_str(vals[i])) for i in range(argc)])
              op_prop = prop_cls(**kwargs)
  
              def infer_shape_entry(num_tensor, tensor_dims,
                                    tensor_shapes, _):
                  """C Callback for ``CustomOpProp::InferShape``."""
                  try:
                      n_in = len(op_prop.list_arguments())
                      n_out = len(op_prop.list_outputs())
                      n_aux = len(op_prop.list_auxiliary_states())
                      assert num_tensor == n_in + n_out + n_aux
  
                      shapes = [[tensor_shapes[i][j] for j in range(tensor_dims[i])]
                                for i in range(n_in)]
                      ret = op_prop.infer_shape(shapes)
                      if len(ret) == 2:
                          ishape, oshape = ret
                          ashape = []
                      elif len(ret) == 3:
                          ishape, oshape, ashape = ret
                      else:
                          raise AssertionError("infer_shape must return 2 or 3 lists")
                      assert len(oshape) == n_out, \
                          "InferShape Error: expecting %d entries in returned output " \
                          "shapes, got %d."%(n_out, len(oshape))
                      assert len(ishape) == n_in, \
                          "InferShape Error: expecting %d entries in returned input " \
                          "shapes, got %d."%(n_in, len(ishape))
                      assert len(ashape) == n_aux, \
                          "InferShape Error: expecting %d entries in returned aux state " \
                          "shapes, got %d."%(n_aux, len(ashape))
                      rshape = list(ishape) + list(oshape) + list(ashape)
                      for i in range(n_in+n_out+n_aux):
                          tensor_shapes[i] = cast(c_array_buf(mx_int,
                                                              array('i', rshape[i])),
                                                  POINTER(mx_int))
                          tensor_dims[i] = len(rshape[i])
  
                      infer_shape_entry._ref_holder = [tensor_shapes]
                  except Exception:
                      print('Error in %s.infer_shape: %s' % (reg_name, traceback.format_exc()))
                      return False
                  return True
  
  
              def infer_storage_type_backward_entry(num_tensor, tensor_stypes, tags, _):
                  # pylint: disable=C0301
                  """C Callback for CustomOpProp::InferStorageTypeBackward"""
                  try:
                      tensors = [[] for i in range(5)]
                      for i in range(num_tensor):
                          tensors[tags[i]].append(_STORAGE_TYPE_ID_TO_STR[tensor_stypes[i]])
                      # Ordering of stypes: ograd, input, output, igrad, aux
                      tensors = [tensors[3], tensors[0], tensors[1], tensors[2], tensors[4]]
                      ret = op_prop.infer_storage_type_backward(tensors[0],
                                                                tensors[1],
                                                                tensors[2],
                                                                tensors[3],
                                                                tensors[4])
                      if len(ret) == 4:
                          ret += []
                      elif len(ret) == 5:
                          pass
                      else:
                          raise AssertionError("infer_storage_type_backward must return 4 or 5 lists")
                      assert len(ret[0]) == len(tensors[0]), \
                          "InferStorageTypeBackward Error: expecting == %d " \
                          "entries in returned output gradient " \
                          "stypes, got %d."%(len(tensors[0]), len(ret[0]))
                      assert len(ret[1]) == len(tensors[1]), \
                          "InferStorageTypeBackward Error: expecting == %d " \
                          "entries in returned input stypes, " \
                          "got %d."%(len(tensors[1]), len(ret[1]))
                      assert len(ret[2]) == len(tensors[2]), \
                          "InferStorageTypeBackward Error: expecting == %d " \
                          "entries in returned output stypes, " \
                          "got %d."%(len(tensors[2]), len(ret[2]))
                      assert len(ret[3]) == len(tensors[3]), \
                          "InferStorageTypeBackward Error: expecting == %d " \
                          "entries in returned input gradient stypes, " \
                          "got %d."%(len(tensors[3]), len(ret[3]))
                      assert len(ret[4]) == len(tensors[4]), \
                          "InferStorageTypeBackward Error: expecting == %d " \
                          "entries in returned aux stypes, " \
                          "got %d."%(len(tensors[4]), len(ret[4]))
                      rstype = []
                      for i, ret_list in enumerate(ret):
                          rstype.extend(ret_list)
  
                      for i, stype in enumerate(rstype):
                          assert stype != _STORAGE_TYPE_ID_TO_STR[_STORAGE_TYPE_UNDEFINED], \
                              "stype should not be undefined"
                          assert stype in _STORAGE_TYPE_STR_TO_ID, \
                              "Provided stype: %s is not valid " \
                              "valid stypes are %s, %s, %s"%(stype,
                                                             _STORAGE_TYPE_ID_TO_STR[_STORAGE_TYPE_DEFAULT],
                                                             _STORAGE_TYPE_ID_TO_STR[_STORAGE_TYPE_ROW_SPARSE],
                                                             _STORAGE_TYPE_ID_TO_STR[_STORAGE_TYPE_CSR])
                          tensor_stypes[i] = _STORAGE_TYPE_STR_TO_ID[stype]
  
                      infer_storage_type_backward_entry._ref_holder = [tensor_stypes]
                  except Exception:
                      print('Error in %s.infer_type: %s' % (reg_name, traceback.format_exc()))
                      return False
                  return True
  
              def infer_storage_type_entry(num_tensor, tensor_stypes, _):
                  """C Callback for CustomOpProp::InferStorageType"""
                  try:
                      n_in = len(op_prop.list_arguments())
                      n_out = len(op_prop.list_outputs())
                      n_aux = len(op_prop.list_auxiliary_states())
                      assert num_tensor == n_in + n_out + n_aux
  
                      stypes = [_STORAGE_TYPE_ID_TO_STR[tensor_stypes[i]] for i in range(n_in)]
                      ret = op_prop.infer_storage_type(stypes)
                      if len(ret) == 2:
                          istype, ostype = ret
                          astype = []
                      elif len(ret) == 3:
                          istype, ostype, astype = ret
                      else:
                          raise AssertionError("infer_storage_type must return 2 or 3 lists")
  
                      assert len(ostype) == n_out, \
                          "InferStorageType Error: expecting %d entries in returned output " \
                          "stypes, got %d."%(n_out, len(ostype))
                      assert len(istype) == n_in, \
                          "InferStorageType Error: expecting %d entries in returned input " \
                          "stypes, got %d."%(n_in, len(istype))
                      assert len(astype) == n_aux, \
                          "InferStorageType Error: expecting %d entries in returned aux state " \
                          "stypes, got %d."%(n_aux, len(astype))
                      rtype = list(istype) + list(ostype) + list(astype)
                      for i, dtype in enumerate(rtype):
                          tensor_stypes[i] = _STORAGE_TYPE_STR_TO_ID[dtype]
                      infer_storage_type_entry._ref_holder = [tensor_stypes]
                  except Exception:
                      print('Error in %s.infer_type: %s' % (reg_name, traceback.format_exc()))
                      return False
                  return True
  
              def infer_type_entry(num_tensor, tensor_types, _):
                  """C Callback for CustomOpProp::InferType"""
                  try:
                      n_in = len(op_prop.list_arguments())
                      n_out = len(op_prop.list_outputs())
                      n_aux = len(op_prop.list_auxiliary_states())
                      assert num_tensor == n_in + n_out + n_aux
  
                      types = [_DTYPE_MX_TO_NP[tensor_types[i]] for i in range(n_in)]
                      ret = op_prop.infer_type(types)
                      if len(ret) == 2:
                          itype, otype = ret
                          atype = []
                      elif len(ret) == 3:
                          itype, otype, atype = ret
                      else:
                          raise AssertionError("infer_type must return 2 or 3 lists")
                      assert len(otype) == n_out, \
                          "InferType Error: expecting %d entries in returned output " \
                          "types, got %d."%(n_out, len(otype))
                      assert len(itype) == n_in, \
                          "InferType Error: expecting %d entries in returned input " \
                          "types, got %d."%(n_in, len(itype))
                      assert len(atype) == n_aux, \
                          "InferType Error: expecting %d entries in returned aux state " \
                          "types, got %d."%(n_aux, len(atype))
                      rtype = list(itype) + list(otype) + list(atype)
                      for i, dtype in enumerate(rtype):
                          tensor_types[i] = _DTYPE_NP_TO_MX[dtype]
  
                      infer_type_entry._ref_holder = [tensor_types]
                  except Exception:
                      print('Error in %s.infer_type: %s' % (reg_name, traceback.format_exc()))
                      return False
                  return True
  
              def list_outputs_entry(out, _):
                  """C Callback for CustomOpProp::ListOutputs"""
                  try:
                      ret = op_prop.list_outputs()
                      ret = [c_str(i) for i in ret] + [c_char_p(0)]
                      ret = c_array(c_char_p, ret)
                      out[0] = cast(ret, POINTER(POINTER(c_char)))
  
                      list_outputs_entry._ref_holder = [out]
                  except Exception:
                      print('Error in %s.list_outputs: %s' % (reg_name, traceback.format_exc()))
                      return False
                  return True
  
              def list_arguments_entry(out, _):
                  """C Callback for CustomOpProp::ListArguments"""
                  try:
                      ret = op_prop.list_arguments()
                      ret = [c_str(i) for i in ret] + [c_char_p(0)]
                      ret = c_array(c_char_p, ret)
                      out[0] = cast(ret, POINTER(POINTER(c_char)))
  
                      list_arguments_entry._ref_holder = [out]
                  except Exception:
                      print('Error in %s.list_arguments: %s' % (reg_name, traceback.format_exc()))
                      return False
                  return True
  
              def list_auxiliary_states_entry(out, _):
                  """C Callback for CustomOpProp::ListAuxiliaryStates"""
                  try:
                      ret = op_prop.list_auxiliary_states()
                      ret = [c_str(i) for i in ret] + [c_char_p(0)]
                      ret = c_array(c_char_p, ret)
                      out[0] = cast(ret, POINTER(POINTER(c_char)))
  
                      list_auxiliary_states_entry._ref_holder = [out]
                  except Exception:
                      tb = traceback.format_exc()
                      print('Error in %s.list_auxiliary_states: %s' % (reg_name, tb))
                      return False
                  return True
  
              def declare_backward_dependency_entry(out_grad, in_data, out_data, num_dep, deps, _):
                  """C Callback for CustomOpProp::DeclareBacwardDependency"""
                  try:
                      out_grad = [out_grad[i] for i in range(len(op_prop.list_outputs()))]
                      in_data = [in_data[i] for i in range(len(op_prop.list_arguments()))]
                      out_data = [out_data[i] for i in range(len(op_prop.list_outputs()))]
                      rdeps = op_prop.declare_backward_dependency(out_grad, in_data, out_data)
                      num_dep[0] = len(rdeps)
                      _registry.result_deps = set()
                      for dep in rdeps:
                          _registry.result_deps.add(dep)
                      rdeps = cast(c_array_buf(c_int, array('i', rdeps)), c_int_p)
                      deps[0] = rdeps
  
                      declare_backward_dependency_entry._ref_holder = [deps]
                  except Exception:
                      tb = traceback.format_exc()
                      print('Error in %s.declare_backward_dependency: %s' % (reg_name, tb))
                      return False
                  return True
  
              def create_operator_entry(ctx, num_inputs, shapes, ndims, dtypes, ret, _):
                  """C Callback for CustomOpProp::CreateOperator"""
                  try:
                      ctx = py_str(ctx)
                      sep = ctx.find('(')
                      ctx = context.Context(ctx[:sep], int(ctx[sep+1:-1]))
                      ndims = [ndims[i] for i in range(num_inputs)]
                      shapes = [[shapes[i][j] for j in range(ndims[i])] for i in range(num_inputs)]
                      dtypes = [dtypes[i] for i in range(num_inputs)]
                      op = op_prop.create_operator(ctx, shapes, dtypes)
  
                      def forward_entry(num_ndarray, ndarraies, tags, reqs, is_train, _):
                          """C Callback for CustomOp::Forward"""
                          try:
                              tensors = [[] for i in range(5)]
                              for i in range(num_ndarray):
                                  if tags[i] == 1 or tags[i] == 4:
                                      tensors[tags[i]].append(_ndarray_cls(cast(ndarraies[i],
                                                                                NDArrayHandle),
                                                                           writable=True))
                                  else:
                                      tensors[tags[i]].append(_ndarray_cls(cast(ndarraies[i],
                                                                                NDArrayHandle),
                                                                           writable=False))
                              reqs = [req_enum[reqs[i]] for i in range(len(tensors[1]))]
                              with ctx:
                                  op.forward(is_train=is_train, req=reqs,
                                             in_data=tensors[0], out_data=tensors[1],
                                             aux=tensors[4])
                          except Exception:
                              print('Error in CustomOp.forward: %s' % traceback.format_exc())
                              return False
                          return True
  
                      def backward_entry(num_ndarray, ndarraies, tags, reqs, is_train, _):
                          """C Callback for CustomOp::Backward"""
                          # pylint: disable=W0613
                          try:
                              tensors = [[] for i in range(5)]
                              num_outputs = len(op_prop.list_outputs())
                              num_args = len(op_prop.list_arguments())
                              for i in range(num_ndarray):
                                  if i in _registry.result_deps or i >= (num_outputs * 2 + num_args):
                                      # If it is a backward dependency or output or aux:
                                      # Set stype as undefined so that it returns
                                      # ndarray based on existing stype
                                      stype = _STORAGE_TYPE_UNDEFINED
                                  else:
                                      # If it is some input, output or out grad ndarray not part of
                                      # backward dependency it is empty and thus the ndarray should
                                      # be set to default
                                      stype = _STORAGE_TYPE_DEFAULT
                                  if tags[i] == 2 or tags[i] == 4:
                                      tensors[tags[i]].append(_ndarray_cls(cast(ndarraies[i],
                                                                                NDArrayHandle),
                                                                           writable=True,
                                                                           stype=stype))
                                  else:
                                      tensors[tags[i]].append(_ndarray_cls(cast(ndarraies[i],
                                                                                NDArrayHandle),
                                                                           writable=False,
                                                                           stype=stype))
                              reqs = [req_enum[reqs[i]] for i in range(len(tensors[2]))]
                              with ctx:
                                  op.backward(req=reqs,
                                              in_data=tensors[0], out_data=tensors[1],
                                              in_grad=tensors[2], out_grad=tensors[3],
                                              aux=tensors[4])
                          except Exception:
                              print('Error in CustomOp.backward: %s' % traceback.format_exc())
                              return False
                          return True
  
                      cur = _registry.inc()
  
                      def delete_entry(_):
                          """C Callback for CustomOp::del"""
                          try:
                              del _registry.ref_holder[cur]
                          except Exception:
                              print('Error in CustomOp.delete: %s' % traceback.format_exc())
                              return False
                          return True
  
                      callbacks = [del_functype(delete_entry),
                                   fb_functype(forward_entry),
                                   fb_functype(backward_entry)]
                      callbacks = [cast(i, CFUNCTYPE(c_int)) for i in callbacks]
                      contexts = [None, None, None]
                      ret[0] = MXCallbackList(c_int(len(callbacks)),
                                              cast(c_array(CFUNCTYPE(c_int), callbacks),
                                                   POINTER(CFUNCTYPE(c_int))),
                                              cast(c_array(c_void_p, contexts),
                                                   POINTER(c_void_p)))
                      op._ref_holder = [ret]
                      _registry.ref_holder[cur] = op
                  except Exception:
                      print('Error in %s.create_operator: %s' % (reg_name, traceback.format_exc()))
                      return False
                  return True
  
              cur = _registry.inc()
  
              def delete_entry(_):
                  """C Callback for CustomOpProp::del"""
                  try:
                      del _registry.ref_holder[cur]
                  except Exception:
                      print('Error in CustomOpProp.delete: %s' % traceback.format_exc())
                      return False
                  return True
  
              callbacks = [del_functype(delete_entry),
                           list_functype(list_arguments_entry),
                           list_functype(list_outputs_entry),
                           list_functype(list_auxiliary_states_entry),
                           infershape_functype(infer_shape_entry),
                           deps_functype(declare_backward_dependency_entry),
                           createop_functype(create_operator_entry),
                           infertype_functype(infer_type_entry),
                           inferstorage_functype(infer_storage_type_entry),
                           inferstorage_backward_functype(infer_storage_type_backward_entry)]
              callbacks = [cast(i, CFUNCTYPE(c_int)) for i in callbacks]
              contexts = [None]*len(callbacks)
              ret[0] = MXCallbackList(c_int(len(callbacks)),
                                      cast(c_array(CFUNCTYPE(c_int), callbacks),
                                           POINTER(CFUNCTYPE(c_int))),
                                      cast(c_array(c_void_p, contexts),
                                           POINTER(c_void_p)))
              op_prop._ref_holder = [ret]
              _registry.ref_holder[cur] = op_prop
              return True
  
          creator_functype = CFUNCTYPE(c_int, c_char_p, c_int, POINTER(c_char_p),
                                       POINTER(c_char_p), POINTER(MXCallbackList))
          creator_func = creator_functype(creator)
          check_call(_LIB.MXCustomOpRegister(c_str(reg_name), creator_func))
          cur = _registry.inc()
          _registry.ref_holder[cur] = creator_func
          return prop_cls
      return do_register
function_tokens: def register ( reg_name ) : def do_register ( prop_cls ) : """Register a subclass of CustomOpProp to the registry.""" fb_functype = CFUNCTYPE ( c_int , c_int , POINTER ( c_void_p ) , POINTER ( c_int ) , POINTER ( c_int ) , c_int , c_void_p ) del_functype = CFUNCTYPE ( c_int , c_void_p ) infershape_functype = CFUNCTYPE ( c_int , c_int , POINTER ( c_int ) , POINTER ( POINTER ( mx_int ) ) , c_void_p ) infertype_functype = CFUNCTYPE ( c_int , c_int , POINTER ( c_int ) , c_void_p ) inferstorage_functype = CFUNCTYPE ( c_int , c_int , POINTER ( c_int ) , c_void_p ) inferstorage_backward_functype = CFUNCTYPE ( c_int , c_int , POINTER ( c_int ) , POINTER ( c_int ) , c_void_p ) list_functype = CFUNCTYPE ( c_int , POINTER ( POINTER ( POINTER ( c_char ) ) ) , c_void_p ) deps_functype = CFUNCTYPE ( c_int , c_int_p , c_int_p , c_int_p , c_int_p , POINTER ( c_int_p ) , c_void_p ) createop_functype = CFUNCTYPE ( c_int , c_char_p , c_int , POINTER ( POINTER ( mx_uint ) ) , POINTER ( c_int ) , POINTER ( c_int ) , POINTER ( MXCallbackList ) , c_void_p ) req_enum = ( 'null' , 'write' , 'inplace' , 'add' ) def creator ( op_type , argc , keys , vals , ret ) : """internal function""" assert py_str ( op_type ) == reg_name kwargs = dict ( [ ( py_str ( keys [ i ] ) , py_str ( vals [ i ] ) ) for i in range ( argc ) ] ) op_prop = prop_cls ( * * kwargs ) def infer_shape_entry ( num_tensor , tensor_dims , tensor_shapes , _ ) : """C Callback for ``CustomOpProp::InferShape``.""" try : n_in = len ( op_prop . list_arguments ( ) ) n_out = len ( op_prop . list_outputs ( ) ) n_aux = len ( op_prop . list_auxiliary_states ( ) ) assert num_tensor == n_in + n_out + n_aux shapes = [ [ tensor_shapes [ i ] [ j ] for j in range ( tensor_dims [ i ] ) ] for i in range ( n_in ) ] ret = op_prop . infer_shape ( shapes ) if len ( ret ) == 2 : ishape , oshape = ret ashape = [ ] elif len ( ret ) == 3 : ishape , oshape , ashape = ret else : raise AssertionError ( "infer_shape must return 2 or 3 lists" ) assert len ( oshape ) == n_out , "InferShape Error: expecting %d entries in returned output " "shapes, got %d." % ( n_out , len ( oshape ) ) assert len ( ishape ) == n_in , "InferShape Error: expecting %d entries in returned input " "shapes, got %d." % ( n_in , len ( ishape ) ) assert len ( ashape ) == n_aux , "InferShape Error: expecting %d entries in returned aux state " "shapes, got %d." % ( n_aux , len ( ashape ) ) rshape = list ( ishape ) + list ( oshape ) + list ( ashape ) for i in range ( n_in + n_out + n_aux ) : tensor_shapes [ i ] = cast ( c_array_buf ( mx_int , array ( 'i' , rshape [ i ] ) ) , POINTER ( mx_int ) ) tensor_dims [ i ] = len ( rshape [ i ] ) infer_shape_entry . _ref_holder = [ tensor_shapes ] except Exception : print ( 'Error in %s.infer_shape: %s' % ( reg_name , traceback . format_exc ( ) ) ) return False return True def infer_storage_type_backward_entry ( num_tensor , tensor_stypes , tags , _ ) : # pylint: disable=C0301 """C Callback for CustomOpProp::InferStorageTypeBackward""" try : tensors = [ [ ] for i in range ( 5 ) ] for i in range ( num_tensor ) : tensors [ tags [ i ] ] . append ( _STORAGE_TYPE_ID_TO_STR [ tensor_stypes [ i ] ] ) # Ordering of stypes: ograd, input, output, igrad, aux tensors = [ tensors [ 3 ] , tensors [ 0 ] , tensors [ 1 ] , tensors [ 2 ] , tensors [ 4 ] ] ret = op_prop . infer_storage_type_backward ( tensors [ 0 ] , tensors [ 1 ] , tensors [ 2 ] , tensors [ 3 ] , tensors [ 4 ] ) if len ( ret ) == 4 : ret += [ ] elif len ( ret ) == 5 : pass else : raise AssertionError ( "infer_storage_type_backward must return 4 or 5 lists" ) assert len ( ret [ 0 ] ) == len ( tensors [ 0 ] ) , "InferStorageTypeBackward Error: expecting == %d " "entries in returned output gradient " "stypes, got %d." % ( len ( tensors [ 0 ] ) , len ( ret [ 0 ] ) ) assert len ( ret [ 1 ] ) == len ( tensors [ 1 ] ) , "InferStorageTypeBackward Error: expecting == %d " "entries in returned input stypes, " "got %d." % ( len ( tensors [ 1 ] ) , len ( ret [ 1 ] ) ) assert len ( ret [ 2 ] ) == len ( tensors [ 2 ] ) , "InferStorageTypeBackward Error: expecting == %d " "entries in returned output stypes, " "got %d." % ( len ( tensors [ 2 ] ) , len ( ret [ 2 ] ) ) assert len ( ret [ 3 ] ) == len ( tensors [ 3 ] ) , "InferStorageTypeBackward Error: expecting == %d " "entries in returned input gradient stypes, " "got %d." % ( len ( tensors [ 3 ] ) , len ( ret [ 3 ] ) ) assert len ( ret [ 4 ] ) == len ( tensors [ 4 ] ) , "InferStorageTypeBackward Error: expecting == %d " "entries in returned aux stypes, " "got %d." % ( len ( tensors [ 4 ] ) , len ( ret [ 4 ] ) ) rstype = [ ] for i , ret_list in enumerate ( ret ) : rstype . extend ( ret_list ) for i , stype in enumerate ( rstype ) : assert stype != _STORAGE_TYPE_ID_TO_STR [ _STORAGE_TYPE_UNDEFINED ] , "stype should not be undefined" assert stype in _STORAGE_TYPE_STR_TO_ID , "Provided stype: %s is not valid " "valid stypes are %s, %s, %s" % ( stype , _STORAGE_TYPE_ID_TO_STR [ _STORAGE_TYPE_DEFAULT ] , _STORAGE_TYPE_ID_TO_STR [ _STORAGE_TYPE_ROW_SPARSE ] , _STORAGE_TYPE_ID_TO_STR [ _STORAGE_TYPE_CSR ] ) tensor_stypes [ i ] = _STORAGE_TYPE_STR_TO_ID [ stype ] infer_storage_type_backward_entry . _ref_holder = [ tensor_stypes ] except Exception : print ( 'Error in %s.infer_type: %s' % ( reg_name , traceback . format_exc ( ) ) ) return False return True def infer_storage_type_entry ( num_tensor , tensor_stypes , _ ) : """C Callback for CustomOpProp::InferStorageType""" try : n_in = len ( op_prop . list_arguments ( ) ) n_out = len ( op_prop . list_outputs ( ) ) n_aux = len ( op_prop . list_auxiliary_states ( ) ) assert num_tensor == n_in + n_out + n_aux stypes = [ _STORAGE_TYPE_ID_TO_STR [ tensor_stypes [ i ] ] for i in range ( n_in ) ] ret = op_prop . infer_storage_type ( stypes ) if len ( ret ) == 2 : istype , ostype = ret astype = [ ] elif len ( ret ) == 3 : istype , ostype , astype = ret else : raise AssertionError ( "infer_storage_type must return 2 or 3 lists" ) assert len ( ostype ) == n_out , "InferStorageType Error: expecting %d entries in returned output " "stypes, got %d." % ( n_out , len ( ostype ) ) assert len ( istype ) == n_in , "InferStorageType Error: expecting %d entries in returned input " "stypes, got %d." % ( n_in , len ( istype ) ) assert len ( astype ) == n_aux , "InferStorageType Error: expecting %d entries in returned aux state " "stypes, got %d." % ( n_aux , len ( astype ) ) rtype = list ( istype ) + list ( ostype ) + list ( astype ) for i , dtype in enumerate ( rtype ) : tensor_stypes [ i ] = _STORAGE_TYPE_STR_TO_ID [ dtype ] infer_storage_type_entry . _ref_holder = [ tensor_stypes ] except Exception : print ( 'Error in %s.infer_type: %s' % ( reg_name , traceback . format_exc ( ) ) ) return False return True def infer_type_entry ( num_tensor , tensor_types , _ ) : """C Callback for CustomOpProp::InferType""" try : n_in = len ( op_prop . list_arguments ( ) ) n_out = len ( op_prop . list_outputs ( ) ) n_aux = len ( op_prop . list_auxiliary_states ( ) ) assert num_tensor == n_in + n_out + n_aux types = [ _DTYPE_MX_TO_NP [ tensor_types [ i ] ] for i in range ( n_in ) ] ret = op_prop . infer_type ( types ) if len ( ret ) == 2 : itype , otype = ret atype = [ ] elif len ( ret ) == 3 : itype , otype , atype = ret else : raise AssertionError ( "infer_type must return 2 or 3 lists" ) assert len ( otype ) == n_out , "InferType Error: expecting %d entries in returned output " "types, got %d." % ( n_out , len ( otype ) ) assert len ( itype ) == n_in , "InferType Error: expecting %d entries in returned input " "types, got %d." % ( n_in , len ( itype ) ) assert len ( atype ) == n_aux , "InferType Error: expecting %d entries in returned aux state " "types, got %d." % ( n_aux , len ( atype ) ) rtype = list ( itype ) + list ( otype ) + list ( atype ) for i , dtype in enumerate ( rtype ) : tensor_types [ i ] = _DTYPE_NP_TO_MX [ dtype ] infer_type_entry . _ref_holder = [ tensor_types ] except Exception : print ( 'Error in %s.infer_type: %s' % ( reg_name , traceback . format_exc ( ) ) ) return False return True def list_outputs_entry ( out , _ ) : """C Callback for CustomOpProp::ListOutputs""" try : ret = op_prop . list_outputs ( ) ret = [ c_str ( i ) for i in ret ] + [ c_char_p ( 0 ) ] ret = c_array ( c_char_p , ret ) out [ 0 ] = cast ( ret , POINTER ( POINTER ( c_char ) ) ) list_outputs_entry . _ref_holder = [ out ] except Exception : print ( 'Error in %s.list_outputs: %s' % ( reg_name , traceback . format_exc ( ) ) ) return False return True def list_arguments_entry ( out , _ ) : """C Callback for CustomOpProp::ListArguments""" try : ret = op_prop . list_arguments ( ) ret = [ c_str ( i ) for i in ret ] + [ c_char_p ( 0 ) ] ret = c_array ( c_char_p , ret ) out [ 0 ] = cast ( ret , POINTER ( POINTER ( c_char ) ) ) list_arguments_entry . _ref_holder = [ out ] except Exception : print ( 'Error in %s.list_arguments: %s' % ( reg_name , traceback . format_exc ( ) ) ) return False return True def list_auxiliary_states_entry ( out , _ ) : """C Callback for CustomOpProp::ListAuxiliaryStates""" try : ret = op_prop . list_auxiliary_states ( ) ret = [ c_str ( i ) for i in ret ] + [ c_char_p ( 0 ) ] ret = c_array ( c_char_p , ret ) out [ 0 ] = cast ( ret , POINTER ( POINTER ( c_char ) ) ) list_auxiliary_states_entry . _ref_holder = [ out ] except Exception : tb = traceback . format_exc ( ) print ( 'Error in %s.list_auxiliary_states: %s' % ( reg_name , tb ) ) return False return True def declare_backward_dependency_entry ( out_grad , in_data , out_data , num_dep , deps , _ ) : """C Callback for CustomOpProp::DeclareBacwardDependency""" try : out_grad = [ out_grad [ i ] for i in range ( len ( op_prop . list_outputs ( ) ) ) ] in_data = [ in_data [ i ] for i in range ( len ( op_prop . list_arguments ( ) ) ) ] out_data = [ out_data [ i ] for i in range ( len ( op_prop . list_outputs ( ) ) ) ] rdeps = op_prop . declare_backward_dependency ( out_grad , in_data , out_data ) num_dep [ 0 ] = len ( rdeps ) _registry . result_deps = set ( ) for dep in rdeps : _registry . result_deps . add ( dep ) rdeps = cast ( c_array_buf ( c_int , array ( 'i' , rdeps ) ) , c_int_p ) deps [ 0 ] = rdeps declare_backward_dependency_entry . _ref_holder = [ deps ] except Exception : tb = traceback . format_exc ( ) print ( 'Error in %s.declare_backward_dependency: %s' % ( reg_name , tb ) ) return False return True def create_operator_entry ( ctx , num_inputs , shapes , ndims , dtypes , ret , _ ) : """C Callback for CustomOpProp::CreateOperator""" try : ctx = py_str ( ctx ) sep = ctx . find ( '(' ) ctx = context . Context ( ctx [ : sep ] , int ( ctx [ sep + 1 : - 1 ] ) ) ndims = [ ndims [ i ] for i in range ( num_inputs ) ] shapes = [ [ shapes [ i ] [ j ] for j in range ( ndims [ i ] ) ] for i in range ( num_inputs ) ] dtypes = [ dtypes [ i ] for i in range ( num_inputs ) ] op = op_prop . create_operator ( ctx , shapes , dtypes ) def forward_entry ( num_ndarray , ndarraies , tags , reqs , is_train , _ ) : """C Callback for CustomOp::Forward""" try : tensors = [ [ ] for i in range ( 5 ) ] for i in range ( num_ndarray ) : if tags [ i ] == 1 or tags [ i ] == 4 : tensors [ tags [ i ] ] . append ( _ndarray_cls ( cast ( ndarraies [ i ] , NDArrayHandle ) , writable = True ) ) else : tensors [ tags [ i ] ] . append ( _ndarray_cls ( cast ( ndarraies [ i ] , NDArrayHandle ) , writable = False ) ) reqs = [ req_enum [ reqs [ i ] ] for i in range ( len ( tensors [ 1 ] ) ) ] with ctx : op . forward ( is_train = is_train , req = reqs , in_data = tensors [ 0 ] , out_data = tensors [ 1 ] , aux = tensors [ 4 ] ) except Exception : print ( 'Error in CustomOp.forward: %s' % traceback . format_exc ( ) ) return False return True def backward_entry ( num_ndarray , ndarraies , tags , reqs , is_train , _ ) : """C Callback for CustomOp::Backward""" # pylint: disable=W0613 try : tensors = [ [ ] for i in range ( 5 ) ] num_outputs = len ( op_prop . list_outputs ( ) ) num_args = len ( op_prop . list_arguments ( ) ) for i in range ( num_ndarray ) : if i in _registry . result_deps or i >= ( num_outputs * 2 + num_args ) : # If it is a backward dependency or output or aux: # Set stype as undefined so that it returns # ndarray based on existing stype stype = _STORAGE_TYPE_UNDEFINED else : # If it is some input, output or out grad ndarray not part of # backward dependency it is empty and thus the ndarray should # be set to default stype = _STORAGE_TYPE_DEFAULT if tags [ i ] == 2 or tags [ i ] == 4 : tensors [ tags [ i ] ] . append ( _ndarray_cls ( cast ( ndarraies [ i ] , NDArrayHandle ) , writable = True , stype = stype ) ) else : tensors [ tags [ i ] ] . append ( _ndarray_cls ( cast ( ndarraies [ i ] , NDArrayHandle ) , writable = False , stype = stype ) ) reqs = [ req_enum [ reqs [ i ] ] for i in range ( len ( tensors [ 2 ] ) ) ] with ctx : op . backward ( req = reqs , in_data = tensors [ 0 ] , out_data = tensors [ 1 ] , in_grad = tensors [ 2 ] , out_grad = tensors [ 3 ] , aux = tensors [ 4 ] ) except Exception : print ( 'Error in CustomOp.backward: %s' % traceback . format_exc ( ) ) return False return True cur = _registry . inc ( ) def delete_entry ( _ ) : """C Callback for CustomOp::del""" try : del _registry . ref_holder [ cur ] except Exception : print ( 'Error in CustomOp.delete: %s' % traceback . format_exc ( ) ) return False return True callbacks = [ del_functype ( delete_entry ) , fb_functype ( forward_entry ) , fb_functype ( backward_entry ) ] callbacks = [ cast ( i , CFUNCTYPE ( c_int ) ) for i in callbacks ] contexts = [ None , None , None ] ret [ 0 ] = MXCallbackList ( c_int ( len ( callbacks ) ) , cast ( c_array ( CFUNCTYPE ( c_int ) , callbacks ) , POINTER ( CFUNCTYPE ( c_int ) ) ) , cast ( c_array ( c_void_p , contexts ) , POINTER ( c_void_p ) ) ) op . _ref_holder = [ ret ] _registry . ref_holder [ cur ] = op except Exception : print ( 'Error in %s.create_operator: %s' % ( reg_name , traceback . format_exc ( ) ) ) return False return True cur = _registry . inc ( ) def delete_entry ( _ ) : """C Callback for CustomOpProp::del""" try : del _registry . ref_holder [ cur ] except Exception : print ( 'Error in CustomOpProp.delete: %s' % traceback . format_exc ( ) ) return False return True callbacks = [ del_functype ( delete_entry ) , list_functype ( list_arguments_entry ) , list_functype ( list_outputs_entry ) , list_functype ( list_auxiliary_states_entry ) , infershape_functype ( infer_shape_entry ) , deps_functype ( declare_backward_dependency_entry ) , createop_functype ( create_operator_entry ) , infertype_functype ( infer_type_entry ) , inferstorage_functype ( infer_storage_type_entry ) , inferstorage_backward_functype ( infer_storage_type_backward_entry ) ] callbacks = [ cast ( i , CFUNCTYPE ( c_int ) ) for i in callbacks ] contexts = [ None ] * len ( callbacks ) ret [ 0 ] = MXCallbackList ( c_int ( len ( callbacks ) ) , cast ( c_array ( CFUNCTYPE ( c_int ) , callbacks ) , POINTER ( CFUNCTYPE ( c_int ) ) ) , cast ( c_array ( c_void_p , contexts ) , POINTER ( c_void_p ) ) ) op_prop . _ref_holder = [ ret ] _registry . ref_holder [ cur ] = op_prop return True creator_functype = CFUNCTYPE ( c_int , c_char_p , c_int , POINTER ( c_char_p ) , POINTER ( c_char_p ) , POINTER ( MXCallbackList ) ) creator_func = creator_functype ( creator ) check_call ( _LIB . MXCustomOpRegister ( c_str ( reg_name ) , creator_func ) ) cur = _registry . inc ( ) _registry . ref_holder [ cur ] = creator_func return prop_cls return do_register