grain.experimental.device_put not exposed

[DBraun]

Using grain 0.2.12 on Linux and macOS:

>>> import grain
>>> grain.__version__
'0.2.12'
>>> grain.experimental.device_put
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
AttributeError: module 'grain.experimental' has no attribute 'device_put'
>>>

So the example usage of grain.experimental.device_put at https://google-grain.readthedocs.io/en/latest/tutorials/dataset_advanced_tutorial.html#:~:text=.experimental.-,device_put,-(%0A%20%20%20%20%20%20%20%20ds%3D is not possible.

One must do

from grain._src.python.experimental.device_put.device_put import device_put
device_put(
    ds=ds,
    device=sharding,
    cpu_buffer_size=cpu_buffer_size,
    device_buffer_size=tpu_buffer_size,
)

[iindyk]

thanks for flagging, that's a bug with internal->external change export, should be fixed by #1006

[sirfz]

Same for grain.experimental.pick_performance_config, I believe it was supposed to be exposed in 0.2.12

[ishanrajsingh]

@selamw1 @iindyk Please review my PR - Fix: Expose device_put in grain.experimental module #1018
#1018