Merge pull request #330 from gomlx/examples-gomod

Fixed gpt2 go.mod to allow for remote execution.

Updated gomlx_checkpoints internal dependency.

Merge pull request #257 from gomlx/dtypes

GoMLX now has its own `dtypes` package, under `pkg/core/dtypes`.

Merge pull request #256 from gomlx/tests

Added GH_TOKEN to GitHub CI tests.

Updated to use new go-xla

Merge branch 'main' of github.com:gomlx/gomlx

Merge pull request #250 from gomlx/v025rc1

Release candidate 1 fixes

Merge pull request #246 from gomlx/distributed

Hightlights:

- Distributed (cross-devices) execution: with AutoSharding and SPMD strategies; 
  Also added support for "portable device" execution.

- API changes: (will require simple fixes)
  - Most not graph building APIs now return errors (as opposed to panicking). Graph building functions
    still use panic to return error -- otherwise it's too painful to express math.
  - All "Rng" renamed to "RNG" -- acronyms in Go are usually capitalized.

- Package `graph`:
  - Fixed/improved documentation.
  - Added `IsNegative`, `IsPositive`, `IsNonNegative`, `IsNonPositive`.
  - Added `SubScalar` and tests for the '*Scalar' functions.
  - Added `Graph.WithDistributedStrategy`, `Graph.WithDeviceMesh`. `Graph.DeviceMesh` and `Graph.NumDevices`
  - Added `Graph.Distributed()` with "collective" (across devices) operations (like `AllReduce`).
  - Renamed: s/`Exec.InDevice`/`Exec.WithDevice`; s/`Exec.SetName`/`Exec.WithName`
  - Added `RunOnDevice`.
  - Added `Exec.AutoSharding` and `Exec.SPMD`.
- Package `context`:
  - Added `context.MustGetParam[T](ctx, key)` and `context.MustGetGraphParam[T](ctx, graph, key)`.
  - Added `Exec.AutoSharding` and `Exec.SPMD`.
  - Added `Variable.DistributedValue` and `Variable.SetDistributedValue`.
- Package `train`:
  - Added `train.DistributedDataset` and `train.BaseDataset`.
  - `Dataset.Reset` now returns an error.
  - `Trainer.TrainStep`, `Trainer.EvalStep` and `Trainer.Eval` now return errors as opposed to panicking.
  - Added `Trainer.WithDeviceAssignment`.
  - Added `Trainer.DistributedTrainStep`, `Trainer.DistributedEvalStep` and `Trainer.DistributedEval`.
- Package `datasets`:
  - Added `datasets.DistributedAccumulator`: converts a normal `Dataset` into a `DistributedDataset`.
  - Added `datasets.OnDevice`: pre-uploads data to devices.
- Package `backend`:
  - Added `Backend.CopyToDevice`
  - `Builder.Parameter()` now takes an optional `ShardingSpec` for sharded inputs.
  - Added ops: `AllReduce`
  - `Backend.NumDevices()` returns an int now.
  - Package `backends/notimplemented`:
    - Added dummy `Backend` that can be used to easily mock backends.
- Package `pkg/core/distributed`:
  -Added `DeviceMesh`, `ShardSpec` and `distributed.Tensor` objects.
- Package `pkg/core/tensors`:
  - Added `Tensor.CheckValid()`, `Tensor.Device()`, `Tensor.Backend()`
  - Changing it to return an error (as opposed to panic) where possible.
Distributed (across devices) execution

Updated README.

Fixed Dockerfile for v0.24.1