def update_stats(self, loss, predictions_dict):
    'Update the accumulators with stats from this batch.\n\n    Args:\n      loss: A scalar tensor representing average loss of examples.\n      predictions_dict: Dictionary of Rank 2 `Tensor` representing information\n          about predictions per example.\n\n    Returns:\n      Three values:\n        - An op that adds a new tree to the ensemble, and\n        - An op that increments the stamp but removes all the trees and resets\n            the handlers. This can be used to reset the state of the ensemble.\n        - A dict containing the training state.\n\n    Raises:\n      ValueError: if inputs are not valid.\n    '
    input_deps = ((self._dense_floats + self._sparse_float_indices) + self._sparse_int_indices)
    worker_device = input_deps[0].device
    predictions = predictions_dict[PREDICTIONS]
    partition_ids = predictions_dict[PARTITION_IDS]
    ensemble_stamp = predictions_dict[ENSEMBLE_STAMP]
    gradients = gradients_impl.gradients(loss, predictions, name='Gradients', colocate_gradients_with_ops=False, gate_gradients=0, aggregation_method=None)[0]
    strategy = self._learner_config.multi_class_strategy
    class_id = self._get_class_id(predictions_dict)
    if (strategy == learner_pb2.LearnerConfig.TREE_PER_CLASS):
        if (self._logits_dimension == 1):
            hessians = gradients_impl.gradients(gradients, predictions, name='Hessian', colocate_gradients_with_ops=False, gate_gradients=0, aggregation_method=None)[0]
            squeezed_gradients = array_ops.squeeze(gradients, axis=[1])
            squeezed_hessians = array_ops.squeeze(hessians, axis=[1])
        else:
            hessian_list = self._diagonal_hessian(gradients, predictions)
            hessians = array_ops.stack(hessian_list, axis=1)
            squeezed_gradients = array_ops.squeeze(_get_column_by_index(gradients, class_id))
            squeezed_hessians = array_ops.squeeze(_get_column_by_index(hessians, class_id))
    else:
        if (strategy == learner_pb2.LearnerConfig.FULL_HESSIAN):
            hessian_list = self._full_hessian(gradients, predictions)
        else:
            hessian_list = self._diagonal_hessian(gradients, predictions)
        squeezed_gradients = gradients
        hessians = array_ops.stack(hessian_list, axis=1)
        squeezed_hessians = hessians
    weights = self._get_weights(self._hessian_shape, squeezed_hessians)
    fc_name_idx = 0
    handlers = []
    init_stamp_token = constant_op.constant(0, dtype=dtypes.int64)
    l1_regularization = constant_op.constant(self._learner_config.regularization.l1, dtypes.float32)
    l2_regularization = constant_op.constant(self._learner_config.regularization.l2, dtypes.float32)
    tree_complexity_regularization = constant_op.constant(self._learner_config.regularization.tree_complexity, dtypes.float32)
    min_node_weight = constant_op.constant(self._learner_config.constraints.min_node_weight, dtypes.float32)
    loss_uses_sum_reduction = (self._loss_reduction == losses.Reduction.SUM)
    loss_uses_sum_reduction = constant_op.constant(loss_uses_sum_reduction)
    weak_learner_type = constant_op.constant(self._learner_config.weak_learner_type)
    num_quantiles = self._num_quantiles
    epsilon = (1.0 / num_quantiles)
    strategy_tensor = constant_op.constant(strategy)
    with ops.device(self._get_replica_device_setter(worker_device)):
        for dense_float_column_idx in range(len(self._dense_floats)):
            fc_name = self._fc_names[fc_name_idx]
            handlers.append(ordinal_split_handler.DenseSplitHandler(l1_regularization=l1_regularization, l2_regularization=l2_regularization, tree_complexity_regularization=tree_complexity_regularization, min_node_weight=min_node_weight, feature_column_group_id=constant_op.constant(dense_float_column_idx), epsilon=epsilon, num_quantiles=num_quantiles, dense_float_column=self._dense_floats[dense_float_column_idx], name=fc_name, gradient_shape=self._gradient_shape, hessian_shape=self._hessian_shape, multiclass_strategy=strategy_tensor, init_stamp_token=init_stamp_token, loss_uses_sum_reduction=loss_uses_sum_reduction, weak_learner_type=weak_learner_type))
            fc_name_idx += 1
        for sparse_float_column_idx in range(len(self._sparse_float_indices)):
            fc_name = self._fc_names[fc_name_idx]
            handlers.append(ordinal_split_handler.SparseSplitHandler(l1_regularization=l1_regularization, l2_regularization=l2_regularization, tree_complexity_regularization=tree_complexity_regularization, min_node_weight=min_node_weight, feature_column_group_id=constant_op.constant(sparse_float_column_idx), epsilon=epsilon, num_quantiles=num_quantiles, sparse_float_column=sparse_tensor.SparseTensor(self._sparse_float_indices[sparse_float_column_idx], self._sparse_float_values[sparse_float_column_idx], self._sparse_float_shapes[sparse_float_column_idx]), name=fc_name, gradient_shape=self._gradient_shape, hessian_shape=self._hessian_shape, multiclass_strategy=strategy_tensor, init_stamp_token=init_stamp_token, loss_uses_sum_reduction=loss_uses_sum_reduction))
            fc_name_idx += 1
        for sparse_int_column_idx in range(len(self._sparse_int_indices)):
            fc_name = self._fc_names[fc_name_idx]
            handlers.append(categorical_split_handler.EqualitySplitHandler(l1_regularization=l1_regularization, l2_regularization=l2_regularization, tree_complexity_regularization=tree_complexity_regularization, min_node_weight=min_node_weight, feature_column_group_id=constant_op.constant(sparse_int_column_idx), sparse_int_column=sparse_tensor.SparseTensor(self._sparse_int_indices[sparse_int_column_idx], self._sparse_int_values[sparse_int_column_idx], self._sparse_int_shapes[sparse_int_column_idx]), name=fc_name, gradient_shape=self._gradient_shape, hessian_shape=self._hessian_shape, multiclass_strategy=strategy_tensor, init_stamp_token=init_stamp_token, loss_uses_sum_reduction=loss_uses_sum_reduction, weak_learner_type=weak_learner_type))
            fc_name_idx += 1
        num_layer_examples = variables.VariableV1(initial_value=array_ops.zeros([], dtypes.int64), name='num_layer_examples', trainable=False)
        num_layer_steps = variables.VariableV1(initial_value=array_ops.zeros([], dtypes.int64), name='num_layer_steps', trainable=False)
        num_layers = variables.VariableV1(initial_value=array_ops.zeros([], dtypes.int64), name='num_layers', trainable=False)
        active_tree = variables.VariableV1(initial_value=array_ops.zeros([], dtypes.int64), name='active_tree', trainable=False)
        active_layer = variables.VariableV1(initial_value=array_ops.zeros([], dtypes.int64), name='active_layer', trainable=False)
        continue_centering = variables.VariableV1(initial_value=self._center_bias, name='continue_centering', trainable=False)
        bias_stats_accumulator = stats_accumulator_ops.StatsAccumulator(stamp_token=0, gradient_shape=self._gradient_shape, hessian_shape=self._hessian_shape, name='BiasAccumulator')
        steps_accumulator = stats_accumulator_ops.StatsAccumulator(stamp_token=0, gradient_shape=tensor_shape.scalar(), hessian_shape=tensor_shape.scalar(), name='StepsAccumulator')
    summary.scalar('layer_stats/num_examples', num_layer_examples)
    summary.scalar('layer_stats/num_steps', num_layer_steps)
    summary.scalar('ensemble_stats/active_tree', active_tree)
    summary.scalar('ensemble_stats/active_layer', active_layer)
    stats_update_ops = []
    stats_update_ops.append(control_flow_ops.cond(continue_centering, self._make_update_bias_stats_fn(ensemble_stamp, predictions, gradients, bias_stats_accumulator), control_flow_ops.no_op))
    handler_reads = collections.OrderedDict()
    for handler in handlers:
        handler_reads[handler] = handler.scheduled_reads()
    handler_results = batch_ops_utils.run_handler_scheduled_ops(handler_reads, ensemble_stamp, worker_device)
    per_handler_updates = collections.OrderedDict()
    subsampling_type = self._learner_config.WhichOneof('feature_fraction')
    if (subsampling_type == 'feature_fraction_per_level'):
        seed = predictions_dict[NUM_LAYERS_ATTEMPTED]
        active_handlers_current_layer = stateless.stateless_random_uniform(shape=[len(handlers)], seed=[seed, 1])
        active_handlers_next_layer = stateless.stateless_random_uniform(shape=[len(handlers)], seed=[(seed + 1), 1])
        active_handlers = array_ops.stack([active_handlers_current_layer, active_handlers_next_layer], axis=1)
        active_handlers = (active_handlers < self._learner_config.feature_fraction_per_level)
    elif (subsampling_type == 'feature_fraction_per_tree'):
        seed = predictions_dict[NUM_TREES_ATTEMPTED]
        active_handlers_current_layer = stateless.stateless_random_uniform(shape=[len(handlers)], seed=[seed, 2])
        active_handlers_current_layer = (active_handlers_current_layer < self._learner_config.feature_fraction_per_tree)
        active_handlers = array_ops.stack([active_handlers_current_layer, array_ops.ones([len(handlers)], dtype=dtypes.bool)], axis=1)
    else:
        active_handlers = array_ops.ones([len(handlers), 2], dtype=dtypes.bool)
    if self._learner_config.constraints.max_number_of_unique_feature_columns:
        target = self._learner_config.constraints.max_number_of_unique_feature_columns

        def _feature_selection_active_handlers():
            used_handlers = array_ops.reshape(predictions_dict[USED_HANDLERS_MASK], [(- 1), 1])
            used_handlers = array_ops.concat([used_handlers, used_handlers], axis=1)
            return math_ops.logical_and(used_handlers, active_handlers)
        active_handlers = control_flow_ops.cond((predictions_dict[NUM_USED_HANDLERS] >= target), _feature_selection_active_handlers, (lambda : active_handlers))
    empty_hess_shape = ([1] + self._hessian_shape.as_list())
    empty_grad_shape = ([1] + self._gradient_shape.as_list())
    empty_gradients = constant_op.constant([], dtype=dtypes.float32, shape=empty_grad_shape)
    empty_hessians = constant_op.constant([], dtype=dtypes.float32, shape=empty_hess_shape)
    active_handlers = array_ops.unstack(active_handlers, axis=0)
    for handler_idx in range(len(handlers)):
        handler = handlers[handler_idx]
        is_active = active_handlers[handler_idx]
        (updates, scheduled_updates) = handler.update_stats(ensemble_stamp, partition_ids, squeezed_gradients, squeezed_hessians, empty_gradients, empty_hessians, weights, is_active, handler_results[handler])
        stats_update_ops.append(updates)
        per_handler_updates[handler] = scheduled_updates
    update_results = batch_ops_utils.run_handler_scheduled_ops(per_handler_updates, ensemble_stamp, worker_device)
    for update in update_results.values():
        stats_update_ops += update
    training_state = GBDTTrainingState(num_layer_examples=num_layer_examples, num_layer_steps=num_layer_steps, num_layers=num_layers, active_tree=active_tree, active_layer=active_layer, continue_centering=continue_centering, bias_stats_accumulator=bias_stats_accumulator, steps_accumulator=steps_accumulator, handlers=handlers)
    reset_op = control_flow_ops.no_op()
    if self._is_chief:
        (stamp_token, _) = model_ops.tree_ensemble_serialize(self._ensemble_handle)
        next_stamp_token = (stamp_token + 1)
        reset_ops = []
        for handler in handlers:
            reset_ops.append(handler.reset(stamp_token, next_stamp_token))
        if self._center_bias:
            reset_ops.append(bias_stats_accumulator.flush(stamp_token, next_stamp_token))
        reset_ops.append(steps_accumulator.flush(stamp_token, next_stamp_token))
        reset_ops.append(self._finalized_trees.assign(0).op)
        reset_ops.append(self._attempted_trees.assign(0).op)
        reset_ops.append(model_ops.tree_ensemble_deserialize(self._ensemble_handle, stamp_token=next_stamp_token, tree_ensemble_config='', name='reset_gbdt'))
        reset_op = control_flow_ops.group([reset_ops])
    return (stats_update_ops, reset_op, training_state)
Update the accumulators with stats from this batch.

Args:
  loss: A scalar tensor representing average loss of examples.
  predictions_dict: Dictionary of Rank 2 `Tensor` representing information
      about predictions per example.

Returns:
  Three values:
    - An op that adds a new tree to the ensemble, and
    - An op that increments the stamp but removes all the trees and resets
        the handlers. This can be used to reset the state of the ensemble.
    - A dict containing the training state.

Raises:
  ValueError: if inputs are not valid.
def update_stats(self, loss, predictions_dict):
    'Update the accumulators with stats from this batch.\n\n    Args:\n      loss: A scalar tensor representing average loss of examples.\n      predictions_dict: Dictionary of Rank 2 `Tensor` representing information\n          about predictions per example.\n\n    Returns:\n      Three values:\n        - An op that adds a new tree to the ensemble, and\n        - An op that increments the stamp but removes all the trees and resets\n            the handlers. This can be used to reset the state of the ensemble.\n        - A dict containing the training state.\n\n    Raises:\n      ValueError: if inputs are not valid.\n    '
    input_deps = ((self._dense_floats + self._sparse_float_indices) + self._sparse_int_indices)
    worker_device = input_deps[0].device
    predictions = predictions_dict[PREDICTIONS]
    partition_ids = predictions_dict[PARTITION_IDS]
    ensemble_stamp = predictions_dict[ENSEMBLE_STAMP]
    gradients = gradients_impl.gradients(loss, predictions, name='Gradients', colocate_gradients_with_ops=False, gate_gradients=0, aggregation_method=None)[0]
    strategy = self._learner_config.multi_class_strategy
    class_id = self._get_class_id(predictions_dict)
    if (strategy == learner_pb2.LearnerConfig.TREE_PER_CLASS):
        if (self._logits_dimension == 1):
            hessians = gradients_impl.gradients(gradients, predictions, name='Hessian', colocate_gradients_with_ops=False, gate_gradients=0, aggregation_method=None)[0]
            squeezed_gradients = array_ops.squeeze(gradients, axis=[1])
            squeezed_hessians = array_ops.squeeze(hessians, axis=[1])
        else:
            hessian_list = self._diagonal_hessian(gradients, predictions)
            hessians = array_ops.stack(hessian_list, axis=1)
            squeezed_gradients = array_ops.squeeze(_get_column_by_index(gradients, class_id))
            squeezed_hessians = array_ops.squeeze(_get_column_by_index(hessians, class_id))
    else:
        if (strategy == learner_pb2.LearnerConfig.FULL_HESSIAN):
            hessian_list = self._full_hessian(gradients, predictions)
        else:
            hessian_list = self._diagonal_hessian(gradients, predictions)
        squeezed_gradients = gradients
        hessians = array_ops.stack(hessian_list, axis=1)
        squeezed_hessians = hessians
    weights = self._get_weights(self._hessian_shape, squeezed_hessians)
    fc_name_idx = 0
    handlers = []
    init_stamp_token = constant_op.constant(0, dtype=dtypes.int64)
    l1_regularization = constant_op.constant(self._learner_config.regularization.l1, dtypes.float32)
    l2_regularization = constant_op.constant(self._learner_config.regularization.l2, dtypes.float32)
    tree_complexity_regularization = constant_op.constant(self._learner_config.regularization.tree_complexity, dtypes.float32)
    min_node_weight = constant_op.constant(self._learner_config.constraints.min_node_weight, dtypes.float32)
    loss_uses_sum_reduction = (self._loss_reduction == losses.Reduction.SUM)
    loss_uses_sum_reduction = constant_op.constant(loss_uses_sum_reduction)
    weak_learner_type = constant_op.constant(self._learner_config.weak_learner_type)
    num_quantiles = self._num_quantiles
    epsilon = (1.0 / num_quantiles)
    strategy_tensor = constant_op.constant(strategy)
    with ops.device(self._get_replica_device_setter(worker_device)):
        for dense_float_column_idx in range(len(self._dense_floats)):
            fc_name = self._fc_names[fc_name_idx]
            handlers.append(ordinal_split_handler.DenseSplitHandler(l1_regularization=l1_regularization, l2_regularization=l2_regularization, tree_complexity_regularization=tree_complexity_regularization, min_node_weight=min_node_weight, feature_column_group_id=constant_op.constant(dense_float_column_idx), epsilon=epsilon, num_quantiles=num_quantiles, dense_float_column=self._dense_floats[dense_float_column_idx], name=fc_name, gradient_shape=self._gradient_shape, hessian_shape=self._hessian_shape, multiclass_strategy=strategy_tensor, init_stamp_token=init_stamp_token, loss_uses_sum_reduction=loss_uses_sum_reduction, weak_learner_type=weak_learner_type))
            fc_name_idx += 1
        for sparse_float_column_idx in range(len(self._sparse_float_indices)):
            fc_name = self._fc_names[fc_name_idx]
            handlers.append(ordinal_split_handler.SparseSplitHandler(l1_regularization=l1_regularization, l2_regularization=l2_regularization, tree_complexity_regularization=tree_complexity_regularization, min_node_weight=min_node_weight, feature_column_group_id=constant_op.constant(sparse_float_column_idx), epsilon=epsilon, num_quantiles=num_quantiles, sparse_float_column=sparse_tensor.SparseTensor(self._sparse_float_indices[sparse_float_column_idx], self._sparse_float_values[sparse_float_column_idx], self._sparse_float_shapes[sparse_float_column_idx]), name=fc_name, gradient_shape=self._gradient_shape, hessian_shape=self._hessian_shape, multiclass_strategy=strategy_tensor, init_stamp_token=init_stamp_token, loss_uses_sum_reduction=loss_uses_sum_reduction))
            fc_name_idx += 1
        for sparse_int_column_idx in range(len(self._sparse_int_indices)):
            fc_name = self._fc_names[fc_name_idx]
            handlers.append(categorical_split_handler.EqualitySplitHandler(l1_regularization=l1_regularization, l2_regularization=l2_regularization, tree_complexity_regularization=tree_complexity_regularization, min_node_weight=min_node_weight, feature_column_group_id=constant_op.constant(sparse_int_column_idx), sparse_int_column=sparse_tensor.SparseTensor(self._sparse_int_indices[sparse_int_column_idx], self._sparse_int_values[sparse_int_column_idx], self._sparse_int_shapes[sparse_int_column_idx]), name=fc_name, gradient_shape=self._gradient_shape, hessian_shape=self._hessian_shape, multiclass_strategy=strategy_tensor, init_stamp_token=init_stamp_token, loss_uses_sum_reduction=loss_uses_sum_reduction, weak_learner_type=weak_learner_type))
            fc_name_idx += 1
        num_layer_examples = variables.VariableV1(initial_value=array_ops.zeros([], dtypes.int64), name='num_layer_examples', trainable=False)
        num_layer_steps = variables.VariableV1(initial_value=array_ops.zeros([], dtypes.int64), name='num_layer_steps', trainable=False)
        num_layers = variables.VariableV1(initial_value=array_ops.zeros([], dtypes.int64), name='num_layers', trainable=False)
        active_tree = variables.VariableV1(initial_value=array_ops.zeros([], dtypes.int64), name='active_tree', trainable=False)
        active_layer = variables.VariableV1(initial_value=array_ops.zeros([], dtypes.int64), name='active_layer', trainable=False)
        continue_centering = variables.VariableV1(initial_value=self._center_bias, name='continue_centering', trainable=False)
        bias_stats_accumulator = stats_accumulator_ops.StatsAccumulator(stamp_token=0, gradient_shape=self._gradient_shape, hessian_shape=self._hessian_shape, name='BiasAccumulator')
        steps_accumulator = stats_accumulator_ops.StatsAccumulator(stamp_token=0, gradient_shape=tensor_shape.scalar(), hessian_shape=tensor_shape.scalar(), name='StepsAccumulator')
    summary.scalar('layer_stats/num_examples', num_layer_examples)
    summary.scalar('layer_stats/num_steps', num_layer_steps)
    summary.scalar('ensemble_stats/active_tree', active_tree)
    summary.scalar('ensemble_stats/active_layer', active_layer)
    stats_update_ops = []
    stats_update_ops.append(control_flow_ops.cond(continue_centering, self._make_update_bias_stats_fn(ensemble_stamp, predictions, gradients, bias_stats_accumulator), control_flow_ops.no_op))
    handler_reads = collections.OrderedDict()
    for handler in handlers:
        handler_reads[handler] = handler.scheduled_reads()
    handler_results = batch_ops_utils.run_handler_scheduled_ops(handler_reads, ensemble_stamp, worker_device)
    per_handler_updates = collections.OrderedDict()
    subsampling_type = self._learner_config.WhichOneof('feature_fraction')
    if (subsampling_type == 'feature_fraction_per_level'):
        seed = predictions_dict[NUM_LAYERS_ATTEMPTED]
        active_handlers_current_layer = stateless.stateless_random_uniform(shape=[len(handlers)], seed=[seed, 1])
        active_handlers_next_layer = stateless.stateless_random_uniform(shape=[len(handlers)], seed=[(seed + 1), 1])
        active_handlers = array_ops.stack([active_handlers_current_layer, active_handlers_next_layer], axis=1)
        active_handlers = (active_handlers < self._learner_config.feature_fraction_per_level)
    elif (subsampling_type == 'feature_fraction_per_tree'):
        seed = predictions_dict[NUM_TREES_ATTEMPTED]
        active_handlers_current_layer = stateless.stateless_random_uniform(shape=[len(handlers)], seed=[seed, 2])
        active_handlers_current_layer = (active_handlers_current_layer < self._learner_config.feature_fraction_per_tree)
        active_handlers = array_ops.stack([active_handlers_current_layer, array_ops.ones([len(handlers)], dtype=dtypes.bool)], axis=1)
    else:
        active_handlers = array_ops.ones([len(handlers), 2], dtype=dtypes.bool)
    if self._learner_config.constraints.max_number_of_unique_feature_columns:
        target = self._learner_config.constraints.max_number_of_unique_feature_columns

        def _feature_selection_active_handlers():
            used_handlers = array_ops.reshape(predictions_dict[USED_HANDLERS_MASK], [(- 1), 1])
            used_handlers = array_ops.concat([used_handlers, used_handlers], axis=1)
            return math_ops.logical_and(used_handlers, active_handlers)
        active_handlers = control_flow_ops.cond((predictions_dict[NUM_USED_HANDLERS] >= target), _feature_selection_active_handlers, (lambda : active_handlers))
    empty_hess_shape = ([1] + self._hessian_shape.as_list())
    empty_grad_shape = ([1] + self._gradient_shape.as_list())
    empty_gradients = constant_op.constant([], dtype=dtypes.float32, shape=empty_grad_shape)
    empty_hessians = constant_op.constant([], dtype=dtypes.float32, shape=empty_hess_shape)
    active_handlers = array_ops.unstack(active_handlers, axis=0)
    for handler_idx in range(len(handlers)):
        handler = handlers[handler_idx]
        is_active = active_handlers[handler_idx]
        (updates, scheduled_updates) = handler.update_stats(ensemble_stamp, partition_ids, squeezed_gradients, squeezed_hessians, empty_gradients, empty_hessians, weights, is_active, handler_results[handler])
        stats_update_ops.append(updates)
        per_handler_updates[handler] = scheduled_updates
    update_results = batch_ops_utils.run_handler_scheduled_ops(per_handler_updates, ensemble_stamp, worker_device)
    for update in update_results.values():
        stats_update_ops += update
    training_state = GBDTTrainingState(num_layer_examples=num_layer_examples, num_layer_steps=num_layer_steps, num_layers=num_layers, active_tree=active_tree, active_layer=active_layer, continue_centering=continue_centering, bias_stats_accumulator=bias_stats_accumulator, steps_accumulator=steps_accumulator, handlers=handlers)
    reset_op = control_flow_ops.no_op()
    if self._is_chief:
        (stamp_token, _) = model_ops.tree_ensemble_serialize(self._ensemble_handle)
        next_stamp_token = (stamp_token + 1)
        reset_ops = []
        for handler in handlers:
            reset_ops.append(handler.reset(stamp_token, next_stamp_token))
        if self._center_bias:
            reset_ops.append(bias_stats_accumulator.flush(stamp_token, next_stamp_token))
        reset_ops.append(steps_accumulator.flush(stamp_token, next_stamp_token))
        reset_ops.append(self._finalized_trees.assign(0).op)
        reset_ops.append(self._attempted_trees.assign(0).op)
        reset_ops.append(model_ops.tree_ensemble_deserialize(self._ensemble_handle, stamp_token=next_stamp_token, tree_ensemble_config=, name='reset_gbdt'))
        reset_op = control_flow_ops.group([reset_ops])
    return (stats_update_ops, reset_op, training_state)