Ejemplo n.º 1
0
static std::array<int64_t, 2> as_array(at::IntList sizes) {
  JIT_ASSERT(sizes.size() == 2);
  std::array<int64_t, 2> arr;
  arr[0] = sizes[0];
  arr[1] = sizes[1];
  return arr;
}
Ejemplo n.º 2
0
// First iterate over the 'from' tensor sizes. Ignore all leading and trailing
// dimensions that are simply one, since they can be trivially broadcasted.
// When iterating over the dimension sizes (with reduced 'from' tensor),
// starting at the trailing dimension, the dimension sizes must either be equal,
// or one of them does not exist.
//
// Note that this is NOT equivalent to numpy broadcasting semantics, and do
// not represent that generalized broadcasting that Pytorch implements in
// general. Rather, this is Caffe2-style broadcasting.
bool fusibleExpandTo(at::IntList from, at::IntList to) {
  if (from.size() > to.size()) {
    return false;
  }
  ssize_t from_dim_start = 0, from_dim_end = from.size() - 1;
  while (from_dim_start < (ssize_t) from.size() && from[from_dim_start] == 1) {
    from_dim_start++;
  }
  while (from_dim_end > from_dim_start && from[from_dim_end] == 1) {
    from_dim_end--;
  }

  ssize_t f = from_dim_end;
  ssize_t t = to.size() - 1;
  for (; f >= from_dim_start && t >= 0; --f, --t) {
    if (from[f] != to[t]) return false;
  }

  // In the case that the 'to' tensor has leading ones in the same place that
  // the 'from' tensor does, f will be less than from_dim_start rather than
  // strictly equal. E.x.: to := [5, 1, 768] and from := [1, 1, 768]
  return f <= from_dim_start;
}