original_docstring: Convert a dict of data into correct dtypes/shapes according to key
docstring: Convert a dict of data into correct dtypes/shapes according to key
docstring_tokens: Convert a dict of data into correct dtypes / shapes according to key
code: def _process_dict(kwargs, ignore_fields=[]):
      """Convert a dict of data into correct dtypes/shapes according to key"""
      # Deal with _some_ dtype issues
      for k, v in kwargs.items():
          if k in ignore_fields:
              continue
  
          if k in _LONG_FIELDS:
              # Any property used as an index must be long (or byte or bool, but those are not relevant for atomic scale systems)
              # int32 would pass later checks, but is actually disallowed by torch
              kwargs[k] = torch.as_tensor(v, dtype=torch.long)
          elif isinstance(v, np.ndarray):
              if np.issubdtype(v.dtype, np.floating):
                  kwargs[k] = torch.as_tensor(v, dtype=torch.get_default_dtype())
              else:
                  kwargs[k] = torch.as_tensor(v)
          elif np.issubdtype(type(v), np.floating):
              # Force scalars to be tensors with a data dimension
              # This makes them play well with irreps
              kwargs[k] = torch.as_tensor(v, dtype=torch.get_default_dtype())
          elif isinstance(v, torch.Tensor) and len(v.shape) == 0:
              # ^ this tensor is a scalar; we need to give it
              # a data dimension to play nice with irreps
              kwargs[k] = v
  
      if AtomicDataDict.BATCH_KEY in kwargs:
          num_frames = kwargs[AtomicDataDict.BATCH_KEY].max() + 1
      else:
          num_frames = 1
  
      for k, v in kwargs.items():
          if k in ignore_fields:
              continue
  
          if len(v.shape) == 0:
              kwargs[k] = v.unsqueeze(-1)
              v = kwargs[k]
  
          if k in set.union(_NODE_FIELDS, _EDGE_FIELDS) and len(v.shape) == 1:
              kwargs[k] = v.unsqueeze(-1)
              v = kwargs[k]
  
          if (
              k in _NODE_FIELDS
              and AtomicDataDict.POSITIONS_KEY in kwargs
              and v.shape[0] != kwargs[AtomicDataDict.POSITIONS_KEY].shape[0]
          ):
              raise ValueError(
                  f"{k} is a node field but has the wrong dimension {v.shape}"
              )
          elif (
              k in _EDGE_FIELDS
              and AtomicDataDict.EDGE_INDEX_KEY in kwargs
              and v.shape[0] != kwargs[AtomicDataDict.EDGE_INDEX_KEY].shape[1]
          ):
              raise ValueError(
                  f"{k} is a edge field but has the wrong dimension {v.shape}"
              )
          elif k in _GRAPH_FIELDS:
              if num_frames > 1 and v.shape[0] != num_frames:
                  raise ValueError(f"Wrong shape for graph property {k}")
code_tokens: def _process_dict ( kwargs , ignore_fields = [ ] ) : for k , v in kwargs . items ( ) : if k in ignore_fields : continue if k in _LONG_FIELDS : kwargs [ k ] = torch . as_tensor ( v , dtype = torch . long ) elif isinstance ( v , np . ndarray ) : if np . issubdtype ( v . dtype , np . floating ) : kwargs [ k ] = torch . as_tensor ( v , dtype = torch . get_default_dtype ( ) ) else : kwargs [ k ] = torch . as_tensor ( v ) elif np . issubdtype ( type ( v ) , np . floating ) : kwargs [ k ] = torch . as_tensor ( v , dtype = torch . get_default_dtype ( ) ) elif isinstance ( v , torch . Tensor ) and len ( v . shape ) == 0 : kwargs [ k ] = v if AtomicDataDict . BATCH_KEY in kwargs : num_frames = kwargs [ AtomicDataDict . BATCH_KEY ] . max ( ) + 1 else : num_frames = 1 for k , v in kwargs . items ( ) : if k in ignore_fields : continue if len ( v . shape ) == 0 : kwargs [ k ] = v . unsqueeze ( - 1 ) v = kwargs [ k ] if k in set . union ( _NODE_FIELDS , _EDGE_FIELDS ) and len ( v . shape ) == 1 : kwargs [ k ] = v . unsqueeze ( - 1 ) v = kwargs [ k ] if ( k in _NODE_FIELDS and AtomicDataDict . POSITIONS_KEY in kwargs and v . shape [ 0 ] != kwargs [ AtomicDataDict . POSITIONS_KEY ] . shape [ 0 ] ) : raise ValueError ( f"{k} is a node field but has the wrong dimension {v.shape}" ) elif ( k in _EDGE_FIELDS and AtomicDataDict . EDGE_INDEX_KEY in kwargs and v . shape [ 0 ] != kwargs [ AtomicDataDict . EDGE_INDEX_KEY ] . shape [ 1 ] ) : raise ValueError ( f"{k} is a edge field but has the wrong dimension {v.shape}" ) elif k in _GRAPH_FIELDS : if num_frames > 1 and v . shape [ 0 ] != num_frames : raise ValueError ( f"Wrong shape for graph property {k}" )
short_docstring: Convert a dict of data into correct dtypes/shapes according to key
short_docstring_tokens: Convert a dict of data into correct dtypes / shapes according to key
comment: """Convert a dict of data into correct dtypes/shapes according to key"""$SEP$# Deal with _some_ dtype issues$SEP$# Any property used as an index must be long (or byte or bool, but those are not relevant for atomic scale systems)$SEP$# int32 would pass later checks, but is actually disallowed by torch$SEP$# Force scalars to be tensors with a data dimension$SEP$# This makes them play well with irreps$SEP$# ^ this tensor is a scalar; we need to give it$SEP$# a data dimension to play nice with irreps