DistRegBuilder.convert_dtype()

DistRegBuilder.convert_dtype()#

DistRegBuilder.convert_dtype(from_dtype, to_dtype)#

Tries to convert the node values in the graph to the specified data type.

Works for nodes whose value is an array or pytree. Nodes whose value is of another type are silently ignored.

Parameters:
  • from_dtype (str | dtype) – The data type to convert from.

  • to_dtype (str | dtype) – The data type to convert to.

Return type:

GraphBuilder

Returns:

The graph builder.