xla

Simplify gradient collection, since we do not need tensors anymore.

#780

Merged