diff --git a/checkall.js b/checkall.js
index 4186da0..4c94edd 100644
--- a/checkall.js
+++ b/checkall.js
@@ -5,13 +5,9 @@
 (function ($) {
 
   Drupal.behaviors.checkAll = {
-    attach: function(context,settings) {
-      $('.form-checkall:not(.checkall-processed)', context).each(function() { Drupal.checkAll.attach(this,settings); });
+    attach: function() {
+      $('.form-checkall:not(.checkall-processed)').each(function() { Drupal.checkAll.attach(this, Drupal.settings); });
     }
-  }
-
-  Drupal.behaviors.checkAll = function(context) {
-    $('.form-checkall:not(.checkall-processed)', context).each(function() { Drupal.checkAll.attach(this); });
   };
 
   Drupal.checkAll = Drupal.checkAll || {
diff --git a/checkall.module b/checkall.module
index 60e5f28..b34daf7 100644
--- a/checkall.module
+++ b/checkall.module
@@ -74,7 +74,7 @@ function checkall_widget_settings_form($field, $instance) {
 /**
  * After build callback.
  */
-function checkall_form_after_build($elements, &$form_state) {
+function checkall_form_after_build(&$elements, &$form_state) {
   checkall_form_after_build_recursive($elements, $form_state);
   return $elements;
 }
@@ -99,8 +99,8 @@ function checkall_form_after_build_recursive(&$elements, &$form_state, $checkall
         $children_checkall = NULL;
 
         // Inherit checkall attribute if specified in CCK field settings.
-        if ($elements[$key]['#type'] == 'container' && isset($elements['#entity']->type)) {
-          $field = field_info_instance('node', $key, $elements['#entity']->type);
+        if ($elements[$key]['#type'] == 'container' && isset($elements['#entity_type'])) {
+          $field = field_info_instance($elements['#entity_type'], $key, $elements['#bundle']);
           if (!empty($field) && isset($field['settings']) && !empty($field['settings']['checkall'])) {
             $children_checkall = TRUE;
           }
@@ -119,7 +119,7 @@ function checkall_form_after_build_recursive(&$elements, &$form_state, $checkall
 /**
  * Process the checkall extension for the checkboxes element.
  */
-function checkall_element_process($element) {
+function checkall_element_process(&$element) {
   // See if the #checkall flag is present.
   if (!empty($element['#checkall'])) {
     $class = is_string($element['#checkall']) ? trim($element['#checkall']) : '';