def test_unweighted_federated_mean_in_xla_execution_context(self):
@computations.federated_computation(
computation_types.FederatedType(np.float32, placements.CLIENTS))
def comp(x):
return intrinsics.federated_mean(x)
execution_contexts.set_local_execution_context()
self.assertEqual(comp([1.0, 2.0, 3.0]), 2.0)
def test_federated_sum_in_xla_execution_context(self):
@computations.federated_computation(
computation_types.FederatedType(np.int32, placements.CLIENTS))
def comp(x):
return intrinsics.federated_sum(x)
execution_contexts.set_local_execution_context()
self.assertEqual(comp([1, 2, 3]), 6)
def test_set_local_execution_context_and_run_simple_xla_computation(self):
builder = xla_client.XlaBuilder('comp')
xla_client.ops.Parameter(builder, 0,
xla_client.shape_from_pyval(tuple()))
xla_client.ops.Constant(builder, np.int32(10))
xla_comp = builder.build()
comp_type = computation_types.FunctionType(None, np.int32)
comp_pb = xla_serialization.create_xla_tff_computation(
xla_comp, [], comp_type)
ctx_stack = context_stack_impl.context_stack
comp = computation_impl.ComputationImpl(comp_pb, ctx_stack)
execution_contexts.set_local_execution_context()
self.assertEqual(comp(), 10)
high=9,
size=(50, ),
dtype=np.int32)
return collections.OrderedDict([('pixels', pixels),
('labels', labels)])
model_type = collections.OrderedDict([
('weights', computation_types.TensorType(np.float32, (784, 10))),
('bias', computation_types.TensorType(np.float32, (10, )))
])
def loss(model, batch):
y = jax.nn.softmax(
jax.numpy.add(
jax.numpy.matmul(batch['pixels'], model['weights']),
model['bias']))
targets = jax.nn.one_hot(jax.numpy.reshape(batch['labels'], -1),
10)
return -jax.numpy.mean(
jax.numpy.sum(targets * jax.numpy.log(y), axis=1))
trainer = jax_components.build_jax_federated_averaging_process(
batch_type, model_type, loss, step_size=0.001)
trainer.next(trainer.initialize(), [[random_batch()]])
if __name__ == '__main__':
execution_contexts.set_local_execution_context()
absltest.main()