(n_devices=None, random_seed=None)
return (is_chief, host_count, n_devices,
          _init_random_number_generators(random_seed))
Initializes host and device attributes for this trainer.

  Args:
    n_devices: Number of devices this trainer will use. If ``None``, get the
        number from the backend.
    random_seed: Random seed as the starting point for all random numbers used
        by the trainer. If ``None``, calculate one from system time and host id.

  Returns:
    is_chief: True if this trainer has special chief responsibilities.
    host_count: Number of hosts in this computation.
    n_devices: The passed in value of n_devices or a computed default (for this
      host).
    random_seed: The passed in value of random_seed or a computed default.
Initializes host and device attributes for this trainer.
Initializes host and device attributes for this trainer .
def init_host_and_devices(n_devices=None, random_seed=None):
  """Initializes host and device attributes for this trainer.

  Args:
    n_devices: Number of devices this trainer will use. If ``None``, get the
        number from the backend.
    random_seed: Random seed as the starting point for all random numbers used
        by the trainer. If ``None``, calculate one from system time and host id.

  Returns:
    is_chief: True if this trainer has special chief responsibilities.
    host_count: Number of hosts in this computation.
    n_devices: The passed in value of n_devices or a computed default (for this
      host).
    random_seed: The passed in value of random_seed or a computed default.
  """
  if fastmath.is_backend(fastmath.Backend.JAX):
    host_id = jax.process_index()
    host_count = jax.host_count()
  else:
    host_id = 0
    host_count = 1
  is_chief = (host_id == 0)

  logging.info('Initializing hosts and devices: host_id %d, host_count %d, '
               'is_chief %d', host_id, host_count, is_chief)

  device_count = fastmath.local_device_count()
  n_devices = n_devices or device_count
  # TODO(lukaszkaiser): remove this restriction when possible.
  if n_devices != device_count and fastmath.is_backend(fastmath.Backend.JAX):
    raise ValueError('JAX cannot work yet with n_devices != all devices: '
                     '%d != %d' % (n_devices, device_count))

  if random_seed is None and host_count > 1:
    random_seed = int(1e6 * (host_id + time.time())) % 2**31
  return (is_chief, host_count, n_devices,
          _init_random_number_generators(random_seed))
def init_host_and_devices ( n_devices = None , random_seed = None ) : if fastmath . is_backend ( fastmath . Backend . JAX ) : host_id = jax . process_index ( ) host_count = jax . host_count ( ) else : host_id = 0 host_count = 1 is_chief = ( host_id == 0 ) logging . info ( 'Initializing hosts and devices: host_id %d, host_count %d, ' 'is_chief %d' , host_id , host_count , is_chief ) device_count = fastmath . local_device_count ( ) n_devices = n_devices or device_count # TODO(lukaszkaiser): remove this restriction when possible. if n_devices != device_count and fastmath . is_backend ( fastmath . Backend . JAX ) : raise ValueError ( 'JAX cannot work yet with n_devices != all devices: ' '%d != %d' % ( n_devices , device_count ) ) if random_seed is None and host_count > 1 : random_seed = int ( 1e6 * ( host_id + time . time ( ) ) ) % 2 ** 31 return ( is_chief , host_count , n_devices , _init_random_number_generators ( random_seed ) )