def get_location(idx, X):
  """Retrieves the location, i.e., the shard index i, slice index j, and value index k of the idx'th element in the struct.

  For example, given a sharded and sliced data object with two shards:
    [[_, _, _], [_, _, _], [_, #, _]], [...]
  Then, the # would be the 7th element (idx), and it would be at location (shard: 0, slice: 2, value: 1).
  """
  num_examples = 0
  for i in range(len(X)):
    for j in range(len(X[i])):
      new_num_examples = num_examples + len(X[i][j])
      if idx < new_num_examples:
        return i, j, idx - num_examples
      num_examples = new_num_examples