```python
return DistArray.from_localarrays(new_key, distribution=a.distribution)
```

There would have to be some logic to handle the case when `a` is a scalar and `b` is the `DistArray`, but in either case, we know the `distribution` ahead of time here.