message: Find allocations hitting the 'replacement' i/p of dynamic slice update.
diff: @@ -146,7 +146,6 @@ AllocationFinder::FindConsumers(const TensorSource& src,
  return;
  }
  case HloOpcode::kDynamicSlice:
  - case HloOpcode::kDynamicUpdateSlice:
  {
  if (op_index == 0) {
  auto t = std::make_pair(user, op_index);
  @@ -154,6 +153,14 @@ AllocationFinder::FindConsumers(const TensorSource& src,
  }
  break;
  }
  + case HloOpcode::kDynamicUpdateSlice:
  + {
  + if (op_index == 0 || op_index==1) {
  + auto t = std::make_pair(user, op_index);
  + tensor_allocation_map.insert(std::make_pair(src, t));
  + }
  + break;
  + }
  case HloOpcode::kCall:
  {
  HloComputation* comp = user->to_apply();
  