code: def batch_row_ids(data_batch):
      
      item = data_batch.data[0]
      user = data_batch.data[1]
      return {'user_weight': user.astype(np.int64),
              'item_weight': item.astype(np.int64)}
docs: Generate row ids based on the current mini-batch